From 25e96452c39f7b65c47b0df5f75cabb46d005023 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 4 Feb 2026 15:44:14 -0500 Subject: [PATCH 01/81] sketch --- src/finch/__init__.py | 4 +- src/finch/compiled.py | 168 ----- src/finch/compiler.py | 31 + src/finch/einstein.py | 561 ---------------- src/finch/linalg/__init__.py | 3 - src/finch/linalg/_linalg.py | 26 - src/finch/tensor.py | 1219 +--------------------------------- 7 files changed, 61 insertions(+), 1951 deletions(-) delete mode 100644 src/finch/compiled.py create mode 100644 src/finch/compiler.py delete mode 100644 src/finch/einstein.py delete mode 100644 src/finch/linalg/__init__.py delete mode 100644 src/finch/linalg/_linalg.py diff --git a/src/finch/__init__.py b/src/finch/__init__.py index 5658656..c853a41 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -134,7 +134,7 @@ ) from .tensor import ( SparseArray, - Tensor, + FinchJLTensor, acos, acosh, all, @@ -225,7 +225,7 @@ "SparseList", "SparseVBL", "Storage", - "Tensor", + "FinchJLTensor", "__array_namespace_info__", "abs", "acos", diff --git a/src/finch/compiled.py b/src/finch/compiled.py deleted file mode 100644 index af05973..0000000 --- a/src/finch/compiled.py +++ /dev/null @@ -1,168 +0,0 @@ -from __future__ import annotations - -from abc import abstractmethod -from collections.abc import Callable, Iterator -from dataclasses import dataclass -from functools import wraps -from typing import TYPE_CHECKING, Any - -from .julia import jl -from .typing import JuliaObj - -if TYPE_CHECKING: - from .tensor import Tensor - -IterObj = tuple | list | dict | Any - - -def _recurse(x: IterObj, /, *, f: Callable[[Any], Any]) -> IterObj: - if isinstance(x, tuple | list): - return type(x)(_recurse(xi, f=f) for xi in x) - if isinstance(x, dict): - ret = {k: _recurse(v, f=f) for k, v in x.items()} - if type(x) is not dict: - ret = type(x)(ret) - return ret - return f(x) - - -def _recurse_iter(x: IterObj, /) -> Iterator[Any]: - if isinstance(x, tuple | list): - for xi in x: - yield from _recurse_iter(xi) - return - if isinstance(x, dict): - for xi in x.values(): - yield from _recurse_iter(xi) - return - yield x - - -def _to_lazy_tensor(x: Tensor | Any, /) -> Tensor | Any: - from .tensor import Tensor - - return x if not isinstance(x, Tensor) else lazy(x) - - -@dataclass -class _ArgumentIndexer: - _idx: int = 0 - - def index(self, _) -> int: - ret = self._idx - self._idx += 1 - return ret - - -def _recurse_iter_compute(x: IterObj, /, *, compute_kwargs: dict[str, Any]) -> IterObj: - from .tensor import Tensor - - # Make a recursive iterator of indices. - idx_obj = _recurse(x, f=_ArgumentIndexer().index) - jl_computed = [] - py_computed = [] - - # Collect lazy tensors; use placeholder - _placeholder = object() - for xi in _recurse_iter(x): - if isinstance(xi, Tensor) and not xi.is_computed(): - jl_computed.append(xi._obj) - py_computed.append(_placeholder) - else: - py_computed.append(xi) - jl_len = len(jl_computed) - # This doesn't return an iterable of arrays -- only a single array - # for `len(jl_computed) == 1` - jl_computed = jl.Finch.compute(*jl_computed, **compute_kwargs) - if jl_len == 1: - jl_computed = (jl_computed,) - - # Replace placeholders with computed tensors. - jl_computed_iter = iter(jl_computed) - for i in range(len(py_computed)): - if py_computed[i] is _placeholder: - py_computed[i] = Tensor(next(jl_computed_iter)) - - # Replace recursive indices by actual computed objects - return _recurse(idx_obj, f=lambda idx: py_computed[idx]) - - -def compiled(opt=None, *, force_materialization=False, tag: int | None = None): - def inner(func): - @wraps(func) - def wrapper_func(*args, **kwargs): - from .tensor import Tensor - - args = tuple(args) - kwargs = dict(kwargs) - compute_at_end = force_materialization or all( - t.is_computed() - for t in _recurse_iter((args, kwargs)) - if isinstance(t, Tensor) - ) - args = _recurse(args, f=_to_lazy_tensor) - kwargs = _recurse(kwargs, f=_to_lazy_tensor) - result = func(*args, **kwargs) - if not compute_at_end: - return result - compute_kwargs = ( - {"ctx": opt.get_julia_scheduler()} if opt is not None else {} - ) - if tag is not None: - compute_kwargs["tag"] = tag - - return _recurse_iter_compute(result, compute_kwargs=compute_kwargs) - - return wrapper_func - - return inner - - -class AbstractScheduler: - def __init__(self, verbose: bool = False): - self.verbose = verbose - - @abstractmethod - def get_julia_scheduler(self) -> JuliaObj: - pass - - -class GalleyScheduler(AbstractScheduler): - def get_julia_scheduler(self) -> JuliaObj: - return jl.Finch.galley_scheduler(verbose=self.verbose) - - -class DefaultScheduler(AbstractScheduler): - def get_julia_scheduler(self) -> JuliaObj: - return jl.Finch.default_scheduler(verbose=self.verbose) - - -def set_optimizer(opt: AbstractScheduler) -> None: - jl.Finch.set_scheduler_b(opt.get_julia_scheduler()) - - -def lazy(tensor: Tensor) -> Tensor: - from .tensor import Tensor - - if tensor.is_computed(): - return Tensor(jl.Finch.LazyTensor(tensor._obj)) - return tensor - - -def compute( - tensor: Tensor, *, opt: AbstractScheduler | None = None, tag: int = -1 -) -> Tensor: - from .tensor import Tensor - - if not tensor.is_computed(): - if opt is None: - return Tensor(jl.Finch.compute(tensor._obj, tag=tag)) - return Tensor( - jl.Finch.compute( - tensor._obj, - verbose=opt.verbose, - ctx=opt.get_julia_scheduler(), - tag=tag, - ) - ) - return tensor diff --git a/src/finch/compiler.py b/src/finch/compiler.py new file mode 100644 index 0000000..4c95bb1 --- /dev/null +++ b/src/finch/compiler.py @@ -0,0 +1,31 @@ +import finchlite + +from finchlite import ( + Loop, + Variable, + Index, + NotationStatement, + NotationModule, +) + +class FinchJLKernel(finchlite.AssemblyKernel): + def __call__(self, *args: FinchJLTensor...) -> tuple[FinchJLTensor...]: + ... + +class FinchJLLibrary(finchlite.AssemblyLibrary): + kernels: dict[str, FinchJLKernel] + def getattr(self, name: str) -> FinchJLKernel: + return self.kernels[name] + +class FinchJLGenerator: + + def __call__(self, prgm: NotationModule) -> FinchJLLibrary: + match prgm: + case NotationModule(statements=stmts): + ... + +class FinchJLCompiler(finchlite.NotationCompiler): + def __call__(self, prgm:NotationModule) -> finchlite.FinchJLLibrary: + generator = FinchJLGenerator() + jl_code = generator(prgm) + return eval(jl_code) \ No newline at end of file diff --git a/src/finch/einstein.py b/src/finch/einstein.py deleted file mode 100644 index 34281cd..0000000 --- a/src/finch/einstein.py +++ /dev/null @@ -1,561 +0,0 @@ -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import Any - -from lark import Lark, Tree - - -class EinsumExpr(ABC): - @abstractmethod - def get_loops(self) -> set[str]: - pass - - @abstractmethod - def run(self, xp, loops, kwargs): - pass - - -nary_ops = { - "+": "add", - "add": "add", - "-": "subtract", - "sub": "subtract", - "subtract": "subtract", - "*": "multiply", - "mul": "multiply", - "multiply": "multiply", - "/": "divide", - "div": "divide", - "divide": "divide", - "//": "floor_divide", - "fld": "floor_divide", - "floor_divide": "floor_divide", - "%": "remainder", - "mod": "remainder", - "remainder": "remainder", - "**": "power", - "pow": "power", - "power": "power", - "==": "equal", - "eq": "equal", - "equal": "equal", - "!=": "not_equal", - "ne": "not_equal", - "not_equal": "not_equal", - "<": "less", - "lt": "less", - "less": "less", - "<=": "less_equal", - "le": "less_equal", - "less_equal": "less_equal", - ">": "greater", - "gt": "greater", - "greater": "greater", - ">=": "greater_equal", - "ge": "greater_equal", - "greater_equal": "greater_equal", - "&": "bitwise_and", - "bitwise_and": "bitwise_and", - "|": "bitwise_or", - "bitwise_or": "bitwise_or", - "^": "bitwise_xor", - "bitwise_xor": "bitwise_xor", - "<<": "bitwise_left_shift", - "lshift": "bitwise_left_shift", - "bitwise_left_shift": "bitwise_left_shift", - ">>": "bitwise_right_shift", - "rshift": "bitwise_right_shift", - "bitwise_right_shift": "bitwise_right_shift", - "and": "logical_and", - "or": "logical_or", - "not": "logical_not", - "min": "minimum", - "max": "maximum", - "logaddexp": "logaddexp", -} - - -unary_ops = { - "+": "positive", - "pos": "positive", - "positive": "positive", - "-": "negative", - "neg": "negative", - "negative": "negative", - "~": "bitwise_invert", - "invert": "bitwise_invert", - "bitwise_invert": "bitwise_invert", - "not": "logical_not", - "logical_not": "logical_not", - "abs": "absolute", - "absolute": "absolute", - "sqrt": "sqrt", - "exp": "exp", - "log": "log", - "log1p": "log1p", - "log10": "log10", - "log2": "log2", - "sin": "sin", - "cos": "cos", - "tan": "tan", - "sinh": "sinh", - "cosh": "cosh", - "tanh": "tanh", - "asin": "arcsin", - "acos": "arccos", - "atan": "arctan", - "asinh": "arcsinh", - "acosh": "arccosh", - "atanh": "arctanh", -} - - -reduction_ops = { - "+": "sum", - "add": "sum", - "sum": "sum", - "*": "prod", - "mul": "prod", - "prod": "prod", - "and": "all", - "or": "any", - "min": "min", - "max": "max", - "argmin": "argmin", - "argmax": "argmax", - "mean": "mean", - "std": "std", - "var": "var", - "count_nonzero": "count_nonzero", - # "&": "bitwise_and", - # "|": "bitwise_or", - # "^": "bitwise_xor", -} - - -@dataclass -class Access(EinsumExpr): - tns: str - idxs: list[str] - - def get_loops(self) -> set[str]: - return set(self.idxs) - - def run(self, xp, loops, kwargs): - assert len(self.idxs) == len(set(self.idxs)) - perm = [self.idxs.index(idx) for idx in loops if idx in self.idxs] - tns = kwargs[self.tns] - tns = xp.permute_dims(tns, perm) - return xp.expand_dims( - tns, [i for i in range(len(loops)) if loops[i] not in self.idxs] - ) - - -@dataclass -class Literal(EinsumExpr): - value: bool | int | float | complex - - def get_loops(self) -> set[str]: - return set() - - def run(self, xp, loops, kwargs): - # Create a scalar array with the same shape as needed - shape = [1] * len(loops) - return xp.full(shape, self.value, format="dense") - - -@dataclass -class Call(EinsumExpr): - func: str - args: list[EinsumExpr] - - def get_loops(self) -> set[str]: - return set().union(*[arg.get_loops() for arg in self.args]) - - def run(self, xp, loops, kwargs): - if len(self.args) == 1: - func = getattr(xp, unary_ops[self.func]) - else: - func = getattr(xp, nary_ops[self.func]) - vals = [arg.run(xp, loops, kwargs) for arg in self.args] - return func(*vals) - - -@dataclass -class Einsum: - arg: EinsumExpr - op: str | None - tns: str - idxs: list[str] - - def run(self, xp, kwargs): - # This is the main entry point for einsum execution - loops = self.arg.get_loops() - assert set(self.idxs).issubset(loops) - loops = sorted(loops) - arg = self.arg.run(xp, loops, kwargs) - axis = tuple(i for i in range(len(loops)) if loops[i] not in self.idxs) - if self.op is not None: - op = getattr(xp, reduction_ops.get(self.op, None)) - val = op(arg, axis=axis) - else: - assert set(self.idxs) == set(loops) - val = arg - dropped = [idx for idx in loops if idx in self.idxs] - axis = [dropped.index(idx) for idx in self.idxs] - return xp.permute_dims(val, axis) - - -lark_parser = Lark(""" - %import common.CNAME - %import common.SIGNED_INT - %import common.SIGNED_FLOAT - %ignore " " // Disregard spaces in text - - start: increment | assign - increment: access (OP | FUNC_NAME) "=" expr - assign: access "=" expr - - // Python operator precedence (lowest to highest) - expr: or_expr - or_expr: and_expr (OR and_expr)* - and_expr: not_expr (AND not_expr)* - not_expr: NOT not_expr | comparison_expr - comparison_expr: bitwise_or_expr ((EQ | NE | LT | LE | GT | GE) bitwise_or_expr)* - bitwise_or_expr: bitwise_xor_expr (PIPE bitwise_xor_expr)* - bitwise_xor_expr: bitwise_and_expr (CARET bitwise_and_expr)* - bitwise_and_expr: shift_expr (AMPERSAND shift_expr)* - shift_expr: add_expr ((LSHIFT | RSHIFT) add_expr)* - add_expr: mul_expr ((PLUS | MINUS) mul_expr)* - mul_expr: unary_expr ((MUL | DIV | FLOORDIV | MOD) unary_expr)* - unary_expr: (PLUS | MINUS | TILDE) unary_expr | power_expr - power_expr: primary (POW unary_expr)? - primary: call_func | access | literal | "(" expr ")" - - OR: "or" - AND: "and" - NOT: "not" - EQ: "==" - NE: "!=" - LT: "<" - LE: "<=" - GT: ">" - GE: ">=" - PIPE: "|" - CARET: "^" - AMPERSAND: "&" - LSHIFT: "<<" - RSHIFT: ">>" - PLUS: "+" - MINUS: "-" - MUL: "*" - DIV: "/" - FLOORDIV: "//" - MOD: "%" - POW: "**" - TILDE: "~" - - OP: "+" | "-" | "*" | "or" | "and" | "|" | "&" | "^" | "<<" | ">>" - | "//" | "/" | "%" | "**" | ">" | "<" | ">=" | "<=" | "==" | "!=" - - access: TNS "[" (IDX ",")* IDX? "]" - call_func: (FUNC_NAME "(" (expr ",")* expr? ")") - literal: bool_literal | complex_literal | float_literal | int_literal - bool_literal: BOOL - int_literal: SIGNED_INT - float_literal: SIGNED_FLOAT - complex_literal: COMPLEX - - BOOL: "True" | "False" - COMPLEX: (SIGNED_FLOAT | SIGNED_INT) ("j" | "J") - IDX: CNAME - TNS: CNAME - FUNC_NAME: CNAME -""") - - -def _parse_einop_expr(t: Tree) -> EinsumExpr: - match t: - case Tree( - "start" - | "expr" - | "or_expr" - | "and_expr" - | "not_expr" - | "comparison_expr" - | "bitwise_or_expr" - | "bitwise_xor_expr" - | "bitwise_and_expr" - | "shift_expr" - | "add_expr" - | "mul_expr" - | "unary_expr" - | "power_expr" - | "primary" - | "literal", - [child], - ): - return _parse_einop_expr(child) - case Tree( - "or_expr" - | "and_expr" - | "bitwise_or_expr" - | "bitwise_and_expr" - | "bitwise_xor_expr" - | "shift_expr" - | "add_expr" - | "mul_expr", - args, - ) if len(args) > 1: - expr = _parse_einop_expr(args[0]) - for i in range(1, len(args), 2): - arg = _parse_einop_expr(args[i + 1]) - expr = Call(args[i].value, [expr, arg]) # type: ignore[union-attr] - return expr - case Tree("comparison_expr", args) if len(args) > 1: - # Handle Python's comparison chaining: a < b < c becomes (a < b) and (b < c) - left = _parse_einop_expr(args[0]) - right = _parse_einop_expr(args[2]) - expr = Call(args[1].value, [left, right]) # type: ignore[union-attr] - for i in range(2, len(args) - 2, 2): - left = _parse_einop_expr(args[i]) - right = _parse_einop_expr(args[i + 2]) - expr = Call("and", [expr, Call(args[i + 1].value, [left, right])]) # type: ignore[union-attr] - return expr - case Tree("power_expr", args) if len(args) > 1: - left = _parse_einop_expr(args[0]) - right = _parse_einop_expr(args[2]) - return Call(args[1].value, [left, right]) # type: ignore[union-attr] - case Tree("unary_expr" | "not_expr", [op, arg]): - return Call(op.value, [_parse_einop_expr(arg)]) # type: ignore[union-attr] - case Tree("access", [tns, *idxs]): - return Access(tns.value, [idx.value for idx in idxs]) # type: ignore[union-attr] - case Tree("bool_literal", [val]): - return Literal(val.value == "True") # type: ignore[union-attr] - case Tree("int_literal", [val]): - return Literal(int(val.value)) # type: ignore[union-attr] - case Tree("float_literal", [val]): - return Literal(float(val.value)) # type: ignore[union-attr] - case Tree("complex_literal", [val]): - return Literal(complex(val.value)) # type: ignore[union-attr] - case Tree("call_func", [func, *args]): - return Call(func.value, [_parse_einop_expr(arg) for arg in args]) # type: ignore[union-attr] - case _: - raise ValueError(f"Unknown tree structure: {t}") - - -def parse_einop(expr: str) -> Einsum: - tree = lark_parser.parse(expr) - - match tree: - case Tree( - "start", [Tree("increment", [Tree("access", [tns, *idxs]), op, expr_node])] - ): - input_expr = _parse_einop_expr(expr_node) # type: ignore[arg-type] - return Einsum(input_expr, op.value, tns.value, [idx.value for idx in idxs]) # type: ignore[union-attr] - - case Tree("start", [Tree("assign", [Tree("access", [tns, *idxs]), expr_node])]): - input_expr = _parse_einop_expr(expr_node) # type: ignore[arg-type] - return Einsum(input_expr, None, tns.value, [idx.value for idx in idxs]) # type: ignore[union-attr] - - case _: - raise ValueError( - f"Expected top-level assignment or increment, got {tree.data}" - ) - - -def einop_impl(xp, prgm, **kwargs): - """Execute an einsum expression using the specified array framework. - - This function parses and executes einsum-like expressions with extended syntax - that supports various operations beyond traditional Einstein summation notation. - - Args: - xp: Array framework module (e.g., numpy, cupy, or other array library) - that provides the underlying array operations. - prgm (str): Einsum program string specifying the computation. The syntax - supports: - - Assignment: "C[i,j] = A[i,j] + B[j,i]" - - Increment: "C[i,j] += A[i,k] * B[k,j]" - - Reductions: "C[i] += A[i,j]", "C[i] max= A[i,j]", "C[i] &= A[i,j]" - - Arithmetic operations: +, -, *, /, //, %, ** - - Comparison operations: ==, !=, <, <=, >, >= - - Logical operations: and, or, not - - Bitwise operations: &, |, ^, <<, >> - - Function calls and complex expressions with parentheses - - Mathematical functions: abs, sqrt, exp, log, sin, cos, tan, etc. - - Literal values: integers, floats, booleans, and complex numbers - - Python operator precedence and parentheses for grouping - **kwargs: Named arrays referenced in the einsum expression. The keys - should match the tensor names used in the program string. - - Returns: - The result array from executing the einsum expression. - - Examples: - >>> import numpy as np - >>> A = np.random.rand(3, 4) - >>> B = np.random.rand(4, 3) - >>> # Matrix addition with transpose - >>> C = einop_impl(np, "C[i,j] = A[i,j] + B[j,i]", A=A, B=B) - >>> # Matrix multiplication - >>> D = einop_impl(np, "D[i,j] += A[i,k] * B[k,j]", A=A, B=B) - >>> # Min-Plus multiplication with shift - >>> E = einop_impl(np, "E[i] min= A[i,k] + D[k,j] << 1", A=A, D=D) - """ - prgm = parse_einop(prgm) - kwargs = {var: xp.Tensor(tns) for var, tns in kwargs.items()} - res = prgm.run(xp, {var: xp.lazy(tns) for var, tns in kwargs.items()}) - if all(tns.is_computed() for tns in kwargs.values()): - return xp.compute(res) - return res - - -def parse_einsum(*args_) -> tuple[Einsum, dict[str, Any]]: - args = list(args_) - if len(args) < 2: - raise ValueError("Expected at least a subscript string and one operand.") - bc = "none" - if isinstance(args[0], str): - subscripts = args[0] - operands = args[1:] - if subscripts.count("->") > 1: - raise ValueError("Subscripts can only contain one '->' symbol.") - if subscripts.count("->") == 1: - subscripts, output_sub = subscripts.split("->") - output_sub = output_sub.strip() - else: - output_sub = None - input_subs = [s.strip() for s in subscripts.split(",")] - # Check for ellipses in input subscripts - if any("..." in sub for sub in input_subs): - if all(sub.startswith("...") for sub in input_subs): - bc = "prefix" - input_subs = [sub[3:] for sub in input_subs] - if output_sub is not None: - assert output_sub.startswith("...") - output_sub = output_sub[3:] - elif all(sub.endswith("...") for sub in input_subs): - bc = "suffix" - input_subs = [sub[:-3] for sub in input_subs] - if output_sub is not None: - assert output_sub.endswith("...") - output_sub = output_sub[:-3] - else: - raise ValueError( - "Ellipses must be at the start or end of all subscripts." - ) - input_idxs = [list(sub) for sub in input_subs] - output_idxs = None if output_sub is None else list(output_sub) - else: - # Alternative syntax: einsum(operand0, subscript0, operand1, subscript1, ...) - # Check if the last element is the output subscript - if len(args) % 2 == 1: - operands = args[0:-2:2] - input_idxs = args[1::2] - output_idxs = list(args[-1]) - output_idxs = [f"j_{j}" for j in output_idxs] - else: - operands = args[0::2] - input_idxs = args[1::2] - output_idxs = None - input_idxs = [[f"j_{j}" for j in idx] for idx in input_idxs] - if any(Ellipsis in idx for idx in input_idxs): - if all(idx[0] == Ellipsis for idx in input_idxs): - bc = "prefix" - input_idxs = [idx[1:] for idx in input_idxs] - if output_idxs is not None: - assert output_idxs[0] == Ellipsis - output_idxs = output_idxs[1:] - elif all(idx[-1] == Ellipsis for idx in input_idxs): - bc = "suffix" - input_idxs = [idx[:-1] for idx in input_idxs] - if output_idxs is not None: - assert output_idxs[-1] == Ellipsis - output_idxs = output_idxs[:-1] - else: - raise ValueError( - "Ellipses must be at the start or end of all subscripts." - ) - - all_idxs = set().union(*input_idxs) - - if output_idxs is None: - output_idx_set = set() - for idx in all_idxs: - if sum(idx in sub for sub in input_idxs) == 1: - output_idx_set.add(idx) - output_idxs = sorted(output_idx_set) - - def ndim(tns): - if hasattr(tns, "ndim"): - return tns.ndim - return 0 - - if bc == "prefix": - max_ell_len = max( - ndim(op) - len(sub) for op, sub in zip(operands, input_idxs, strict=False) - ) - for i in range(len(operands)): - ell_idxs = [ - f"i_{j}" - for j in range( - max_ell_len - (ndim(operands[i]) - len(input_idxs[i])), max_ell_len - ) - ] - input_idxs[i] = ell_idxs + input_idxs[i] - ell_idxs = [f"i_{j}" for j in range(max_ell_len)] - output_idxs = [f"i_{j}" for j in range(max_ell_len)] + output_idxs - elif bc == "suffix": - max_ell_len = max( - ndim(op) - len(sub) for op, sub in zip(operands, input_idxs, strict=False) - ) - for i in range(len(operands)): - ell_idxs = [f"k_{j}" for j in range(ndim(operands[i]) - len(input_idxs[i]))] - input_idxs[i] = input_idxs[i] + ell_idxs - output_idxs = output_idxs + [f"k_{j}" for j in range(max_ell_len)] - - all_idxs = set().union(*input_idxs) - - if len(input_idxs) != len(operands): - raise ValueError("Number of input subscripts must match number of operands.") - assert set(output_idxs).issubset(all_idxs), ( - "Output indices must be a subset of input indices." - ) - tag = 0 - - def freshen(x): - nonlocal tag - tag += 1 - return f"{x}_{tag}" - - for j in all_idxs: - freshen(j) - op = None if output_idxs == all_idxs else "add" - out_tns = freshen("B") - idxs = tuple(output_idxs) - in_tnss = [freshen("A") for _ in operands] - arg = Access(in_tnss[0], input_idxs[0]) - for i in range(1, len(operands)): - arg = Call( - "mul", - (arg, Access(in_tnss[i], input_idxs[i])), - ) # type: ignore[assignment] - return ( - Einsum( - arg, - op, - out_tns, - idxs, - ), - {in_tnss[i]: operands[i] for i in range(len(operands))}, - ) - - -def einsum_impl(xp, *args): - prgm, kwargs = parse_einsum(*args) - kwargs = {var: xp.Tensor(tns) for var, tns in kwargs.items()} - res = prgm.run(xp, {var: xp.lazy(tns) for var, tns in kwargs.items()}) - if all(tns.is_computed() for tns in kwargs.values()): - return xp.compute(res) - return res diff --git a/src/finch/linalg/__init__.py b/src/finch/linalg/__init__.py deleted file mode 100644 index 55a7867..0000000 --- a/src/finch/linalg/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from ._linalg import vector_norm - -__all__ = ["vector_norm"] diff --git a/src/finch/linalg/_linalg.py b/src/finch/linalg/_linalg.py deleted file mode 100644 index 7c598ce..0000000 --- a/src/finch/linalg/_linalg.py +++ /dev/null @@ -1,26 +0,0 @@ -from numpy.core.numeric import normalize_axis_tuple - -from ..julia import jl -from ..tensor import Tensor - - -def vector_norm( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, - ord: float = 2, -) -> Tensor: - if axis is not None: - axis = normalize_axis_tuple(axis, x.ndim) - if axis != tuple(range(x.ndim)): - raise ValueError( - "At the moment only `None` (vector norm of a flattened array) " - "is supported. Got: {axis}." - ) - - result = Tensor(jl.Finch.norm(x._obj, ord)) - if keepdims: - result = result[tuple(None for _ in range(x.ndim))] - return result diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 7f260ce..5bcf328 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -9,8 +9,6 @@ from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple from . import dtypes as jl_dtypes -from .compiled import compiled, compute, lazy -from .einstein import einop_impl, einsum_impl from .errors import PerformanceWarning from .julia import jc, jl from .levels import ( @@ -24,6 +22,7 @@ sparse_formats_names, ) from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix +from finchlite import Tensor, TensorFType class SparseArray: @@ -31,29 +30,39 @@ class SparseArray: PyData/Sparse marker class """ +class FinchJLTensorFType(TensorFType): + def __init__(self, jltype): + self.jltype = jltype -class Tensor(_Display, SparseArray): + def ndims(self) -> int: + return jl.ndims(self.jltype) + + def element_type(self): + return jl.eltype(self.jltype) + + ... + +class FinchJLTensor(_Display, SparseArray, Tensor, finchlite.EagerTensor): """ A wrapper class for Finch.Tensor and Finch.SwizzleArray. Constructors ------------ - Tensor(scipy.sparse.spmatrix) + FinchJLTensor(scipy.sparse.spmatrix) Construct a Tensor out of a `scipy.sparse` object. Supported formats are: `COO`, `CSC`, and `CSR`. - Tensor(numpy.ndarray) + FinchJLTensor(numpy.ndarray) Construct a Tensor out of a NumPy array object. This is a no-copy operation. - Tensor(Storage) + FinchJLTensor(Storage) Initialize a Tensor with a `storage` description. `storage` can already hold data. - Tensor(julia_object) - Tensor created from a compatible raw Julia object. Must be a `SwizzleArray` or - `LazyTensor`. + FinchJLTensor(julia_object) + Tensor created from a compatible raw Julia object. Must be a `Tensor`. This is a no-copy operation. Parameters ---------- - obj : np.ndarray or scipy.sparse or Storage or Finch.SwizzleArray + obj : np.ndarray or scipy.sparse or Storage or Finch.Tensor Input to construct a Tensor. It's a no-copy operation of for NumPy and SciPy input. For Storage it's levels' description with order. The order numbers the dimensions from the fastest to slowest. The leaf nodes have @@ -69,15 +78,15 @@ class Tensor(_Display, SparseArray): Returns ------- - Tensor - Python wrapper for Finch Tensor. + FinchJLTensor + Python wrapper for Finch.jl `Tensor`. Examples -------- >>> import numpy as np >>> import finch >>> arr2d = np.arange(6).reshape((2, 3)) - >>> t1 = finch.Tensor(arr2d) + >>> t1 = finch.FinchJLTensor(arr2d) >>> t1.todense() array([[0, 1, 2], [3, 4, 5]]) @@ -92,9 +101,6 @@ class Tensor(_Display, SparseArray): [3, 4, 5]]) """ - row_major: str = "C" - column_major: str = "F" - def __init__( self, obj: np.ndarray | spmatrix | Storage | JuliaObj, @@ -129,11 +135,11 @@ def __init__( if copy: self._raise_julia_copy_not_supported() self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1))) - elif jl.isa(obj, jl.Finch.SwizzleArray) or jl.isa(obj, jl.Finch.LazyTensor): + elif jl.isa(obj, jl.Finch.Tensor): if copy: self._raise_julia_copy_not_supported() self._obj = obj - elif isinstance(obj, Tensor): + elif isinstance(obj, FinchJLTensor): self._obj = obj._obj else: raise ValueError( @@ -141,145 +147,9 @@ def __init__( f"be provided. Found: {type(obj)}" ) - def __pos__(self): - return self._elemwise_op("+") - - def __neg__(self): - return self._elemwise_op("-") - - def __add__(self, other): - return self._elemwise_op("+", other) - - def __mul__(self, other): - return self._elemwise_op("*", other) - - def __sub__(self, other): - return self._elemwise_op("-", other) - - def __truediv__(self, other): - return self._elemwise_op("/", other) - - def __floordiv__(self, other): - return self._elemwise_op("Finch.fld_nothrow", other) - - def __mod__(self, other): - return self._elemwise_op("Finch.mod_nothrow", other) - - def __pow__(self, other): - return self._elemwise_op("^", other) - - @compiled() - def __matmul__(self, other: Tensor) -> Tensor: - if self.ndim == 0 or other.ndim == 0: - raise ValueError( - f"`{self.ndim=}`, `{other.ndim=}`. Both must be greater than `0`." - ) - - if other.ndim == 1: - return sum(self * other, axis=-1) - - if self.ndim == 1: - return sum(self * other.mT, axis=-1) - - return sum(self[..., :, None, :] * other.mT[..., None, :, :], axis=-1) - - def __abs__(self): - return self._elemwise_op("abs") - - def __invert__(self): - return self._elemwise_op("~") - - def __and__(self, other): - return self._elemwise_op("&", other) - - def __or__(self, other): - return self._elemwise_op("|", other) - - def __xor__(self, other): - return self._elemwise_op("xor", other) - - def __lshift__(self, other): - return self._elemwise_op("<<", other) - - def __rshift__(self, other): - return self._elemwise_op(">>", other) - - def __lt__(self, other): - return self._elemwise_op("<", other) - - def __le__(self, other): - return self._elemwise_op("<=", other) - - def __gt__(self, other): - return self._elemwise_op(">", other) - - def __ge__(self, other): - return self._elemwise_op(">=", other) - - def __eq__(self, other): - return self._elemwise_op("==", other) - - def __ne__(self, other): - return self._elemwise_op("!=", other) - - def _elemwise_op(self, op: str, other: Tensor | None = None) -> Tensor: - if other is None: - result = jl.broadcast(jl.seval(op), self._obj) - else: - if np.isscalar(other): - other = jc.convert(self.dtype, other) - else: - other = jl.permutedims(other._obj, tuple(range(other.ndim, 0, -1))) - # inverse swizzle, so `broadcast` appends new dims to the front - result = jl.broadcast( - jl.seval(op), - jl.permutedims(self._obj, tuple(range(self.ndim, 0, -1))), - other, - ) - # swizzle back to the original order - result = jl.permutedims(result, tuple(range(jl.ndims(result), 0, -1))) - - return Tensor(result) - - def __bool__(self): - return self._to_scalar(bool) - - def __float__(self): - return self._to_scalar(float) - - def __int__(self): - return self._to_scalar(int) - - def __index__(self): - return self._to_scalar(int) - - def __complex__(self): - return self._to_scalar(complex) - - def _to_scalar(self, builtin): - if self.ndim != 0: - raise ValueError(f"{builtin} can be computed for one-element tensors only.") - return builtin(self.todense().flatten()[0]) - - def __getitem__(self, key): - if not isinstance(key, tuple): - key = (key,) - - if not self.is_computed(): - # lazy indexing mode - key = _process_lazy_indexing(key, self.ndim) - else: - # standard indexing mode - key = _expand_ellipsis(key, self.shape) - key = _add_missing_dims(key, self.shape) - key = _add_plus_one(key, self.shape) - - result = self._obj[key] - if jl.isa(result, jl.Finch.SwizzleArray) or jl.isa(result, jl.Finch.LazyTensor): - return Tensor(result) - if jl.isa(result, jl.Finch.Tensor): - return Tensor(jl.swizzle(result, *range(1, jl.ndims(result) + 1))) - return result + @property + def element_type(self): + return jl.eltype(self._obj.body) @property def dtype(self) -> DType: @@ -333,31 +203,6 @@ def to_device( return self - def is_computed(self) -> bool: - return not jl.isa(self._obj, jl.Finch.LazyTensor) - - @classmethod - def preprocess_order(cls, order: OrderType, ndim: int) -> tuple[int, ...]: - if order == cls.column_major: - permutation = tuple(range(1, ndim + 1)) - elif order == cls.row_major or order is None: - permutation = tuple(range(1, ndim + 1)[::-1]) - elif isinstance(order, tuple): - if builtins.min(order) == 0: - order = tuple(i + 1 for i in order) - if len(order) == ndim and builtins.all( - i in order for i in range(1, ndim + 1) - ): - permutation = order - else: - raise ValueError(f"Custom order is not a permutation: {order}.") - else: - raise ValueError( - f"order must be 'C', 'F' or a tuple, but is: {type(order)}." - ) - - return permutation - @classmethod def get_lvl_ndim(cls, lvl: JuliaObj) -> int: ndim = 0 @@ -368,18 +213,6 @@ def get_lvl_ndim(cls, lvl: JuliaObj) -> int: break return ndim - def get_order(self, zero_indexing: bool = True) -> tuple[int, ...]: - order = self._order - if zero_indexing: - order = tuple(i - 1 for i in order) - return order - - def get_inv_order(self, zero_indexing: bool = True) -> tuple[int, ...]: - inv_order = jl.invperm(self._order) - if zero_indexing: - inv_order = tuple(i - 1 for i in inv_order) - return inv_order - def todense(self) -> np.ndarray: obj = self._obj @@ -401,6 +234,7 @@ def todense(self) -> np.ndarray: result = np.asarray(jl.reshape(dense_tensor.val, shape)) return result.transpose(self.get_order()) if self._is_dense else result + #TODO: Do we need? def permute_dims(self, axes: tuple[int, ...]) -> Tensor: axes = tuple(i + 1 for i in axes) new_obj = jl.permutedims(self._obj, axes) @@ -729,1003 +563,6 @@ def reshape( return Tensor(arr) -def full( - shape: int | tuple[int, ...], - fill_value: jl_dtypes.number, - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - _validate_device(device) - if not np.isscalar(fill_value): - raise ValueError("`fill_value` must be a scalar") - if format not in ("coo", "dense"): - raise ValueError(f"{format} format not supported.") - if isinstance(shape, int): - shape = (shape,) - dtype = ( - np.asarray(fill_value).dtype.type - if dtype is None - else jl_dtypes.jl_to_np_dtype[dtype] - ) - if dtype == np.bool_: # Fails with: Finch currently only supports isbits defaults - dtype = bool - - if format == "coo" and shape != (): - return Tensor( - jl.Tensor(jl.SparseCOO[len(shape)](jl.Element(dtype(fill_value))), *shape) - ) - # for dense format or () shape - return Tensor(np.full(shape, fill_value, dtype=dtype)) - - -def full_like( - x: Tensor, - /, - fill_value: jl_dtypes.number, - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - return full(x.shape, fill_value, dtype=dtype, format=format, device=device) - - -def ones( - shape: int | tuple[int, ...], - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - return full(shape, np.float64(1), dtype=dtype, format=format, device=device) - - -def ones_like( - x: Tensor, - /, - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - dtype = x.dtype if dtype is None else dtype - return ones(x.shape, dtype=dtype, format=format, device=device) - - -def zeros( - shape: int | tuple[int, ...], - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - return full(shape, np.float64(0), dtype=dtype, format=format, device=device) - - -def zeros_like( - x: Tensor, - /, - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - dtype = x.dtype if dtype is None else dtype - return zeros(x.shape, dtype=dtype, format=format, device=device) - - -def empty( - shape: int | tuple[int, ...], - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - return full(shape, np.float64(0), dtype=dtype, format=format, device=device) - - -def empty_like( - x: Tensor, - /, - *, - dtype: DType | None = None, - format: str = "coo", - device: Device = None, -) -> Tensor: - dtype = x.dtype if dtype is None else dtype - return empty(x.shape, dtype=dtype, format=format, device=device) - - -def arange( - start: float, - /, - stop: float | None = None, - step: float = 1, - *, - dtype: DType | None = None, - device: Device = None, -) -> Tensor: - _validate_device(device) - return Tensor(np.arange(start, stop, step, jl_dtypes.jl_to_np_dtype[dtype])) - - -def linspace( - start: complex, - stop: complex, - /, - num: int, - *, - dtype: DType | None = None, - device: Device = None, - endpoint: bool = True, -) -> Tensor: - _validate_device(device) - return Tensor( - np.linspace( - start, - stop, - num=num, - dtype=jl_dtypes.jl_to_np_dtype[dtype], - endpoint=endpoint, - ) - ) - - -def permute_dims(x: Tensor, axes: tuple[int, ...]) -> Tensor: - return x.permute_dims(axes) - - -def moveaxis(x: Tensor, source: int, destination: int) -> Tensor: - axes = list(range(x.ndim)) - norm_source = normalize_axis_index(source, x.ndim) - norm_dest = normalize_axis_index(destination, x.ndim) - axes.insert(norm_dest, axes.pop(norm_source)) - return x.permute_dims(tuple(axes)) - - -def astype(x: Tensor, dtype: DType, /, *, copy: bool = True) -> Tensor: - if not copy: - if x.dtype == dtype: - return x - if copy is False: - raise ValueError("Unable to avoid a copy while casting in no-copy mode.") - - finch_tns = x._obj.body - result = jl.copyto_b( - jl.similar(finch_tns, jc.convert(dtype, jl.fill_value(finch_tns)), dtype), - finch_tns, - ) - return Tensor(jl.swizzle(result, *x.get_order(zero_indexing=False))) - - -def where(condition: Tensor, x1: Tensor, x2: Tensor, /) -> Tensor: - axis_cond, axis_x1, axis_x2 = ( - range(condition.ndim, 0, -1), - range(x1.ndim, 0, -1), - range(x2.ndim, 0, -1), - ) - # inverse swizzle, so `broadcast` appends new dims to the front - result = jl.broadcast( - jl.ifelse, - jl.permutedims(condition._obj, tuple(axis_cond)), - jl.permutedims(x1._obj, tuple(axis_x1)), - jl.permutedims(x2._obj, tuple(axis_x2)), - ) - # swizzle back to the original order - result = jl.permutedims(result, tuple(range(jl.ndims(result), 0, -1))) - return Tensor(result) - - -def nonzero(x: Tensor, /) -> tuple[np.ndarray, ...]: - indices = jl.ffindnz(x._obj)[:-1] # return only indices, skip values - indices = tuple(np.asarray(i) - 1 for i in indices) - sort_order = np.lexsort(indices[::-1]) # sort to row-major, C-style order - return tuple(Tensor(i[sort_order]) for i in indices) - - -def _reduce_core( - x: Tensor, fn: Callable, axis: int | tuple[int, ...] | None, keepdims: bool = False -): - if axis is None: - axis = tuple(range(x.ndim)) - axis = normalize_axis_tuple(axis, x.ndim) - axis = tuple(i + 1 for i in axis) - if keepdims: - return fn(x._obj, dims=axis) - if x.is_computed(): - return jl.compute(jl.dropdims(fn(jl.lazy(x._obj), dims=axis), dims=axis)) - return jl.dropdims(fn(x._obj, dims=axis), dims=axis) - - -def _reduce_sum_prod( - x: Tensor, - fn: Callable, - axis: int | tuple[int, ...] | None, - dtype: DType | None, - keepdims: bool = False, -) -> Tensor: - result = _reduce_core(x, fn, axis, keepdims) - - if np.isscalar(result): - tmp_dtype = jl_dtypes.int_ if jl.seval(f"{x.dtype} <: Integer") else x.dtype - result = jl.Tensor( - jl.Element( - jc.convert(tmp_dtype, 0), - np.array(result, dtype=jl_dtypes.jl_to_np_dtype[tmp_dtype]), - ) - ) - - result = Tensor(result) - - if jl.isa(result._obj, jl.Finch.LazyTensor): - if dtype is not None: - raise ValueError( - "`dtype` keyword for `sum` and `prod` in the lazy mode isn't supported" - ) - # dtype casting rules - elif dtype is not None: - result = astype(result, dtype, copy=None) - elif jl.seval(f"{x.dtype} <: Unsigned"): - result = astype(result, jl_dtypes.uint, copy=None) - elif jl.seval(f"{x.dtype} <: Signed"): - result = astype(result, jl_dtypes.int_, copy=None) - - return result - - -def _reduce( - x: Tensor, fn: Callable, axis: int | tuple[int, ...] | None, keepdims: bool = False -) -> Tensor: - result = _reduce_core(x, fn, axis, keepdims) - if np.isscalar(result): - result = jl.Tensor( - jl.Element( - jc.convert(x.dtype, 0), - np.array(result, dtype=jl_dtypes.jl_to_np_dtype[x.dtype]), - ) - ) - return Tensor(result) - - -def sum( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - dtype: DType | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce_sum_prod(x, jl.sum, axis, dtype, keepdims) - - -def prod( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - dtype: DType | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce_sum_prod(x, jl.prod, axis, dtype, keepdims) - - -def max( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce(x, jl.maximum, axis, keepdims) - - -def min( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce(x, jl.minimum, axis, keepdims) - - -def any( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce(x != 0, jl.any, axis, keepdims) - - -def all( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce(x != 0, jl.all, axis, keepdims) - - -def mean( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce(x, jl.mean, axis, keepdims) - - -def std( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - correction: float = 0.0, - keepdims: bool = False, -) -> Tensor: - def _std(x, dims): - return jl.std(x, correction=correction, dims=dims) - - return _reduce(x, _std, axis, keepdims) - - -def var( - x: Tensor, - /, - *, - axis: int | tuple[int, ...] | None = None, - correction: float = 0.0, - keepdims: bool = False, -) -> Tensor: - def _var(x, dims): - return jl.var(x, correction=correction, dims=dims) - - return _reduce(x, _var, axis, keepdims) - - -def argmin( - x: Tensor, - /, - *, - axis: int | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce(x, jl.Finch.argmin_python, axis, keepdims) - - -def argmax( - x: Tensor, - /, - *, - axis: int | None = None, - keepdims: bool = False, -) -> Tensor: - return _reduce(x, jl.Finch.argmax_python, axis, keepdims) - - -def squeeze( - x: Tensor, - /, - axis: int | tuple[int, ...], -) -> Tensor: - if isinstance(axis, int): - axis = (axis,) - axis = normalize_axis_tuple(axis, x.ndim) - axis = tuple(i + 1 for i in axis) - result = jl.dropdims(x._obj, dims=axis) - return Tensor(result) - - -def expand_dims( - x: Tensor, - /, - axis: int | tuple[int, ...] = 0, -) -> Tensor: - if isinstance(axis, int): - axis = (axis,) - axis = normalize_axis_tuple(axis, x.ndim + len(axis)) - axis = tuple(i + 1 for i in axis) - result = jl.expanddims(x._obj, dims=axis) - return Tensor(result) - - -def diagonal(x: Tensor, /, *, offset: int = 0) -> Tensor: - m = x.shape[-2] - n = x.shape[-1] - mask = eye(m, n, k=offset, format="coo", dtype=bool) - res = compute(sum(where(mask, lazy(x), zeros(x.shape)), axis=-1)) - return res[..., 0 : builtins.min(m, n, m + offset, n - offset)] - - -def eye( - n_rows: int, - n_cols: int | None = None, - /, - *, - k: int = 0, - dtype: DType | None = None, - format: Literal["coo", "dense"] = "coo", - device: Device = None, -) -> Tensor: - _validate_device(device) - n_cols = n_rows if n_cols is None else n_cols - dtype = jl_dtypes.float64 if dtype is None else dtype - tns = jl.Finch.eye_python(n_rows, n_cols, k, dtype(False)) - if format == "coo": - return Tensor(tns) - if format == "dense": - return Tensor(jl.Tensor(jl.DenseFormat(2, dtype(False)), tns)) - raise ValueError(f"{format} not supported, only 'coo' and 'dense' is allowed.") - - -def tensordot(x1: Tensor, x2: Tensor, /, *, axes=2) -> Tensor: - if not isinstance(x1, Tensor): - x1 = Tensor(x1) - if not isinstance(x2, Tensor): - x2 = Tensor(x2) - if isinstance(axes, Iterable): - self_axes = normalize_axis_tuple(axes[0], x1.ndim) - other_axes = normalize_axis_tuple(axes[1], x2.ndim) - axes = (tuple(i + 1 for i in self_axes), tuple(i + 1 for i in other_axes)) - - result = jl.tensordot(x1._obj, x2._obj, axes) - return Tensor(result) - - -def log(x: Tensor, /) -> Tensor: - return x._elemwise_op("log") - - -def log10(x: Tensor, /) -> Tensor: - return x._elemwise_op("log10") - - -def log1p(x: Tensor, /) -> Tensor: - return x._elemwise_op("log1p") - - -def log2(x: Tensor, /) -> Tensor: - return x._elemwise_op("log2") - - -def sqrt(x: Tensor, /) -> Tensor: - return x._elemwise_op("sqrt") - - -def sign(x: Tensor, /) -> Tensor: - return x._elemwise_op("sign") - - -def round(x: Tensor, /) -> Tensor: - return x._elemwise_op("round") - - -def isnan(x: Tensor, /) -> Tensor: - return x._elemwise_op("isnan") - - -def isinf(x: Tensor, /) -> Tensor: - return x._elemwise_op("isinf") - - -def isfinite(x: Tensor, /) -> Tensor: - return x._elemwise_op("isfinite") - - -def exp(x: Tensor, /) -> Tensor: - return x._elemwise_op("exp") - - -def expm1(x: Tensor, /) -> Tensor: - return x._elemwise_op("expm1") - - -def floor(x: Tensor, /) -> Tensor: - return x._elemwise_op("floor") - - -def ceil(x: Tensor, /) -> Tensor: - return x._elemwise_op("ceil") - - -def cos(x: Tensor, /) -> Tensor: - return x._elemwise_op("cos") - - -def cosh(x: Tensor, /) -> Tensor: - return x._elemwise_op("cosh") - - -def acos(x: Tensor, /) -> Tensor: - return x._elemwise_op("acos") - - -def acosh(x: Tensor, /) -> Tensor: - return x._elemwise_op("acosh") - - -def sin(x: Tensor, /) -> Tensor: - return x._elemwise_op("sin") - - -def sinh(x: Tensor, /) -> Tensor: - return x._elemwise_op("sinh") - - -def asin(x: Tensor, /) -> Tensor: - return x._elemwise_op("asin") - - -def asinh(x: Tensor, /) -> Tensor: - return x._elemwise_op("asinh") - - -def tan(x: Tensor, /) -> Tensor: - return x._elemwise_op("tan") - - -def tanh(x: Tensor, /) -> Tensor: - return x._elemwise_op("tanh") - - -def atan(x: Tensor, /) -> Tensor: - return x._elemwise_op("atan") - - -def atanh(x: Tensor, /) -> Tensor: - return x._elemwise_op("atanh") - - -def atan2(x1: Tensor, x2: Tensor, /) -> Tensor: - return x1._elemwise_op("atand", x2) - - -def trunc(x: Tensor, /) -> Tensor: - return x._elemwise_op("trunc") - - -def real(x: Tensor, /) -> Tensor: - return x._elemwise_op("real") - - -def imag(x: Tensor, /) -> Tensor: - return x._elemwise_op("imag") - - -def conj(x: Tensor, /) -> Tensor: - return x._elemwise_op("conj") - - -def square(x: Tensor, /) -> Tensor: - return x ** Tensor(2) - - -def logaddexp(x1: Tensor, x2: Tensor, /) -> Tensor: - return log(exp(x1) + exp(x2)) - - -def logical_and(x1: Tensor, x2: Tensor, /) -> Tensor: - return x1._elemwise_op("Finch.and", x2) - - -def logical_or(x1: Tensor, x2: Tensor, /) -> Tensor: - return x1._elemwise_op("Finch.or", x2) - - -def logical_xor(x1: Tensor, x2: Tensor, /) -> Tensor: - return x1._elemwise_op("Finch.xor", x2) - - -def power(x1: Tensor, x2: Tensor, /) -> Tensor: - return x1._elemwise_op("^", x2) - - -def einop(prgm, **kwargs): - """Execute an einsum-like expression - - This function parses and executes einsum-like expressions with extended syntax - that supports various operations beyond traditional Einstein summation notation. - - Args: - prgm (str): Einsum program string specifying the computation. The syntax - supports: - - Assignment: "C[i,j] = A[i,j] + B[j,i]" - - Increment: "C[i,j] += A[i,k] * B[k,j]" - - Reductions: "C[i] += A[i,j]", "C[i] max= A[i,j]", "C[i] &= A[i,j]" - - Arithmetic operations: +, -, *, /, //, %, ** - - Comparison operations: ==, !=, <, <=, >, >= - - Logical operations: and, or, not - - Bitwise operations: &, |, ^, <<, >> - - Function calls and complex expressions with parentheses - - Mathematical functions: abs, sqrt, exp, log, sin, cos, tan, etc. - - Literal values: integers, floats, booleans, and complex numbers - - Python operator precedence and parentheses for grouping - **kwargs: Named arrays referenced in the einsum expression. The keys - should match the tensor names used in the program string. - - Returns: - The result array from executing the einsum expression. - - Examples: - >>> A = finch.random.rand(3, 4) - >>> B = finch.random.rand(4, 3) - >>> # Matrix addition with transpose - >>> C = finch.einop("C[i,j] = A[i,j] + B[j,i]", A=A, B=B) - >>> # Matrix multiplication - >>> D = finch.einop("D[i,j] += A[i,k] * B[k,j]", A=A, B=B) - >>> # Min-Plus multiplication with shift - >>> E = finch.einop("E[i] min= A[i,k] + D[k,j] << 1", A=A, D=D) - """ - import finch - - return einop_impl(finch, prgm, **kwargs) - - -def einsum(*args, **kwargs): - """ - einsum(subscripts, *operands) - - Evaluates the Einstein summation convention on the operands. - - Using the Einstein summation convention, many common multi-dimensional, - linear algebraic array operations can be represented in a simple fashion. - In *implicit* mode `einsum` computes these values. - - In *explicit* mode, `einsum` provides further flexibility to compute - other array operations that might not be considered classical Einstein - summation operations, by disabling, or forcing summation over specified - subscript labels. - - See the notes and examples for clarification. - - Parameters - ---------- - subscripts : str - Specifies the subscripts for summation as comma separated list of - subscript labels. An implicit (classical Einstein summation) - calculation is performed unless the explicit indicator '->' is - included as well as subscript labels of the precise output form. - operands : list of array_like - These are the arrays for the operation. - - Returns - ------- - output : ndarray - The calculation based on the Einstein summation convention. - - Notes - ----- - The Einstein summation convention can be used to compute - many multi-dimensional, linear algebraic array operations. `einsum` - provides a succinct way of representing these. - - A non-exhaustive list of these operations, - which can be computed by `einsum`, is shown below along with examples: - - * Trace of an array, :py:func:`numpy.trace`. - * Return a diagonal, :py:func:`numpy.diag`. - * Array axis summations, :py:func:`numpy.sum`. - * Transpositions and permutations, :py:func:`numpy.transpose`. - * Matrix multiplication and dot product, :py:func:`numpy.matmul` - :py:func:`numpy.dot`. - * Vector inner and outer products, :py:func:`numpy.inner` - :py:func:`numpy.outer`. - * Broadcasting, element-wise and scalar multiplication, - :py:func:`numpy.multiply`. - * Tensor contractions, :py:func:`numpy.tensordot`. - * Chained array operations, in efficient calculation order, - :py:func:`numpy.einsum_path`. - - The subscripts string is a comma-separated list of subscript labels, - where each label refers to a dimension of the corresponding operand. - Whenever a label is repeated it is summed, so ``np.einsum('i,i', a, b)`` - is equivalent to :py:func:`np.inner(a,b) `. If a label - appears only once, it is not summed, so ``np.einsum('i', a)`` - produces a view of ``a`` with no changes. A further example - ``np.einsum('ij,jk', a, b)`` describes traditional matrix multiplication - and is equivalent to :py:func:`np.matmul(a,b) `. - Repeated subscript labels in one operand take the diagonal. - For example, ``np.einsum('ii', a)`` is equivalent to - :py:func:`np.trace(a) `. - - In *implicit mode*, the chosen subscripts are important - since the axes of the output are reordered alphabetically. This - means that ``np.einsum('ij', a)`` doesn't affect a 2D array, while - ``np.einsum('ji', a)`` takes its transpose. Additionally, - ``np.einsum('ij,jk', a, b)`` returns a matrix multiplication, while, - ``np.einsum('ij,jh', a, b)`` returns the transpose of the - multiplication since subscript 'h' precedes subscript 'i'. - - In *explicit mode* the output can be directly controlled by - specifying output subscript labels. This requires the - identifier '->' as well as the list of output subscript labels. - This feature increases the flexibility of the function since - summing can be disabled or forced when required. The call - ``np.einsum('i->', a)`` is like :py:func:`np.sum(a) ` - if ``a`` is a 1-D array, and ``np.einsum('ii->i', a)`` - is like :py:func:`np.diag(a) ` if ``a`` is a square 2-D array. - The difference is that `einsum` does not allow broadcasting by default. - Additionally ``np.einsum('ij,jh->ih', a, b)`` directly specifies the - order of the output subscript labels and therefore returns matrix - multiplication, unlike the example above in implicit mode. - - To enable and control broadcasting, use an ellipsis. Default - NumPy-style broadcasting is done by adding an ellipsis - to the left of each term, like ``np.einsum('...ii->...i', a)``. - ``np.einsum('...i->...', a)`` is like - :py:func:`np.sum(a, axis=-1) ` for array ``a`` of any shape. - To take the trace along the first and last axes, - you can do ``np.einsum('i...i', a)``, or to do a matrix-matrix - product with the left-most indices instead of rightmost, one can do - ``np.einsum('ij...,jk...->ik...', a, b)``. - - `einsum` also provides an alternative way to provide the subscripts and - operands as ``einsum(op0, sublist0, op1, sublist1, ..., [sublistout])``. - If the output shape is not provided in this format `einsum` will be - calculated in implicit mode, otherwise it will be performed explicitly. - The examples below have corresponding `einsum` calls with the two - parameter methods. - - Examples - -------- - >>> a = np.arange(25).reshape(5, 5) - >>> b = np.arange(5) - >>> c = np.arange(6).reshape(2, 3) - - Trace of a matrix: - - >>> np.einsum("ii", a) - 60 - >>> np.einsum(a, [0, 0]) - 60 - >>> np.trace(a) - 60 - - Extract the diagonal (requires explicit form): - - >>> np.einsum("ii->i", a) - array([ 0, 6, 12, 18, 24]) - >>> np.einsum(a, [0, 0], [0]) - array([ 0, 6, 12, 18, 24]) - >>> np.diag(a) - array([ 0, 6, 12, 18, 24]) - - Sum over an axis (requires explicit form): - - >>> np.einsum("ij->i", a) - array([ 10, 35, 60, 85, 110]) - >>> np.einsum(a, [0, 1], [0]) - array([ 10, 35, 60, 85, 110]) - >>> np.sum(a, axis=1) - array([ 10, 35, 60, 85, 110]) - - For higher dimensional arrays summing a single axis can be done - with ellipsis: - - >>> np.einsum("...j->...", a) - array([ 10, 35, 60, 85, 110]) - >>> np.einsum(a, [Ellipsis, 1], [Ellipsis]) - array([ 10, 35, 60, 85, 110]) - - Compute a matrix transpose, or reorder any number of axes: - - >>> np.einsum("ji", c) - array([[0, 3], - [1, 4], - [2, 5]]) - >>> np.einsum("ij->ji", c) - array([[0, 3], - [1, 4], - [2, 5]]) - >>> np.einsum(c, [1, 0]) - array([[0, 3], - [1, 4], - [2, 5]]) - >>> np.transpose(c) - array([[0, 3], - [1, 4], - [2, 5]]) - - Vector inner products: - - >>> np.einsum("i,i", b, b) - 30 - >>> np.einsum(b, [0], b, [0]) - 30 - >>> np.inner(b, b) - 30 - - Matrix vector multiplication: - - >>> np.einsum("ij,j", a, b) - array([ 30, 80, 130, 180, 230]) - >>> np.einsum(a, [0, 1], b, [1]) - array([ 30, 80, 130, 180, 230]) - >>> np.dot(a, b) - array([ 30, 80, 130, 180, 230]) - >>> np.einsum("...j,j", a, b) - array([ 30, 80, 130, 180, 230]) - - Broadcasting and scalar multiplication: - - >>> np.einsum("..., ...", 3, c) - array([[ 0, 3, 6], - [ 9, 12, 15]]) - >>> np.einsum(",ij", 3, c) - array([[ 0, 3, 6], - [ 9, 12, 15]]) - >>> np.einsum(3, [Ellipsis], c, [Ellipsis]) - array([[ 0, 3, 6], - [ 9, 12, 15]]) - >>> np.multiply(3, c) - array([[ 0, 3, 6], - [ 9, 12, 15]]) - - Vector outer product: - - >>> np.einsum("i,j", np.arange(2) + 1, b) - array([[0, 1, 2, 3, 4], - [0, 2, 4, 6, 8]]) - >>> np.einsum(np.arange(2) + 1, [0], b, [1]) - array([[0, 1, 2, 3, 4], - [0, 2, 4, 6, 8]]) - >>> np.outer(np.arange(2) + 1, b) - array([[0, 1, 2, 3, 4], - [0, 2, 4, 6, 8]]) - - Tensor contraction: - - >>> a = np.arange(60.0).reshape(3, 4, 5) - >>> b = np.arange(24.0).reshape(4, 3, 2) - >>> np.einsum("ijk,jil->kl", a, b) - array([[4400., 4730.], - [4532., 4874.], - [4664., 5018.], - [4796., 5162.], - [4928., 5306.]]) - >>> np.einsum(a, [0, 1, 2], b, [1, 0, 3], [2, 3]) - array([[4400., 4730.], - [4532., 4874.], - [4664., 5018.], - [4796., 5162.], - [4928., 5306.]]) - >>> np.tensordot(a, b, axes=([1, 0], [0, 1])) - array([[4400., 4730.], - [4532., 4874.], - [4664., 5018.], - [4796., 5162.], - [4928., 5306.]]) - - Example of ellipsis use: - - >>> a = np.arange(6).reshape((3, 2)) - >>> b = np.arange(12).reshape((4, 3)) - >>> np.einsum("ki,jk->ij", a, b) - array([[10, 28, 46, 64], - [13, 40, 67, 94]]) - >>> np.einsum("ki,...k->i...", a, b) - array([[10, 28, 46, 64], - [13, 40, 67, 94]]) - >>> np.einsum("k...,jk", a, b) - array([[10, 28, 46, 64], - [13, 40, 67, 94]]) - """ - import finch - - return einsum_impl(finch, *args) - - -def _is_scipy_sparse_obj(x): - return hasattr(x, "__module__") and x.__module__.startswith("scipy.sparse") - - -def _slice_plus_one(s: slice, size: int) -> range: - step = s.step if s.step is not None else 1 - start_default = size if step < 0 else 1 - stop_default = 1 if step < 0 else size - - if s.start is not None: - start = normalize_axis_index(s.start, size) + 1 if s.start < size else size - else: - start = start_default - - if s.stop is not None: - stop_offset = 2 if step < 0 else 0 - stop = ( - normalize_axis_index(s.stop, size) + stop_offset if s.stop < size else size - ) - else: - stop = stop_default - - if (start, stop, step) == (1, size, 1): - return jl.Colon() - - return jl.range(start=start, step=step, stop=stop) - - -def _add_plus_one(key: tuple, shape: tuple[int, ...]) -> tuple: - new_key = [] - sizes = iter(shape) - for idx in key: - if idx is None: - new_key.append(jl.nothing) - continue - - size = next(sizes) - if isinstance(idx, int): - new_key.append(normalize_axis_index(idx, size) + 1) - elif isinstance(idx, slice): - new_key.append(_slice_plus_one(idx, size)) - elif isinstance(idx, list | np.ndarray | tuple): - idx = normalize_axis_tuple(idx, size) - new_key.append(jl.Vector([i + 1 for i in idx])) - else: - new_key.append(idx) - - return tuple(new_key) - - -def _expand_ellipsis(key: tuple, shape: tuple[int, ...]) -> tuple: - ellipsis_pos = None - key_without_ellipsis = [] - # first we need to find the ellipsis and confirm it's the only one - for pos, idx in enumerate(key): - if idx is Ellipsis: - if ellipsis_pos is None: - ellipsis_pos = pos - else: - raise IndexError("an index can only have a single ellipsis ('...')") - else: - key_without_ellipsis.append(idx) - key = key_without_ellipsis - - # then we expand ellipsis with a full range - if ellipsis_pos is not None: - n_missing_idxs = len(shape) - builtins.sum(1 for k in key if k is not None) - key = key[:ellipsis_pos] + [slice(None)] * n_missing_idxs + key[ellipsis_pos:] - - return tuple(key) - - -def _add_missing_dims(key: tuple, shape: tuple[int, ...]) -> tuple: - missing_dims = len(shape) - builtins.sum(1 for k in key if k is not None) - return key + (slice(None),) * missing_dims - - -def _process_lazy_indexing(key: tuple, ndim: int) -> tuple: - new_key = () - ellipsis_found = False - for idx in key: - if idx == slice(None): - new_key += (jl.Colon(),) - elif idx is None: - new_key += (jl.nothing,) - elif idx is Ellipsis: - num_of_colons = ndim - builtins.sum(1 for k in key if k is not None) + 1 - new_key += (jl.Colon(),) * num_of_colons - if ellipsis_found: - raise IndexError("an index can only have a single ellipsis ('...')") - - ellipsis_found = True - else: - raise ValueError(f"Invalid lazy index member: {idx}") - return new_key - - -def _eq_scalars(x, y): - if x is None or y is None: - return x == y - if jl.isnan(x) or jl.isnan(y): - return jl.isnan(x) and jl.isnan(y) - return x == y - - def _validate_device(device: Device) -> None: if device not in {"cpu", None}: raise ValueError( From 48b6b902307f54459b8846a7823ee41add868347 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 10 Feb 2026 20:32:43 -0500 Subject: [PATCH 02/81] feat: improved the compiler interface --- pyproject.toml | 3 ++- src/finch/compiler.py | 57 ++++++++++++++++++++++++++----------------- 2 files changed, 36 insertions(+), 24 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 999b30c..e0ec2fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,13 +4,14 @@ version = "0.2.14" description = "Sparse Tensor Programming in Python powered by Finch.jl" readme = "README.md" authors = [{name = "Willow Ahrens", email = "willow.marie.ahrens@gmail.com"}] -requires-python = ">=3.11,<4.0" +requires-python = ">=3.11,<3.14" dependencies = [ "juliapkg (>=0.1.16,<0.2.0)", "numpy (>=1.19,<2.4)", "numba>=0.61,<0.63.1", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", + "finch-tensor-lite==0.2.0" ] [tool.poetry] diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 4c95bb1..d0bbee6 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,31 +1,42 @@ -import finchlite - -from finchlite import ( - Loop, - Variable, - Index, - NotationStatement, - NotationModule, -) - -class FinchJLKernel(finchlite.AssemblyKernel): - def __call__(self, *args: FinchJLTensor...) -> tuple[FinchJLTensor...]: - ... - -class FinchJLLibrary(finchlite.AssemblyLibrary): - kernels: dict[str, FinchJLKernel] - def getattr(self, name: str) -> FinchJLKernel: +from .tensor import FinchJLTensor +from finchlite import NotationCompiler, AssemblyKernel, AssemblyLibrary +import finchlite.finch_notation.nodes as ntn + +from juliacall import Main as jl + +CompiledModuleName = "compiled_module" +FunctionName = "compiled_function" + +class FinchJLKernel(AssemblyKernel): + def __init__(self, func_name, jl_code): + self.func_name = func_name + jl.seval(jl_code) + def __call__( + self, *args: tuple[FinchJLTensor, ...] + ) -> tuple[FinchJLTensor, ...]: + argList = [] + for arg in args: + argList.append(f"arg{len(argList)}") + setattr(jl, argList[-1], arg) + + jl.seval(f"{self.func_name}({argList.join(',')})") + +class FinchJLLibrary(AssemblyLibrary): + def __init__(self, kernel_name, kernel): + self.kernel_dict = {kernel_name: kernel} + + def __getattr__(self, name: str) -> FinchJLKernel: return self.kernels[name] -class FinchJLGenerator: - def __call__(self, prgm: NotationModule) -> FinchJLLibrary: +class FinchJLGenerator: + def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: match prgm: - case NotationModule(statements=stmts): + case ntn.Module(statements=stmts): ... -class FinchJLCompiler(finchlite.NotationCompiler): - def __call__(self, prgm:NotationModule) -> finchlite.FinchJLLibrary: +class FinchJLCompiler(NotationCompiler): + def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: generator = FinchJLGenerator() jl_code = generator(prgm) - return eval(jl_code) \ No newline at end of file + return FinchJLLibrary(CompiledModuleName, FinchJLKernel(FunctionName, jl_code)) \ No newline at end of file From 9ed21968bdb62500d0d36682008acf5091ef7284 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 10 Feb 2026 22:47:57 -0500 Subject: [PATCH 03/81] wip: preparing to test compiler --- src/finch/__init__.py | 713 ++++++++++----------- src/finch/_array_api_info.py | 174 +++--- src/finch/compiler.py | 87 ++- src/finch/dtypes.py | 122 ++-- src/finch/errors.py | 4 +- src/finch/io.py | 20 +- src/finch/levels.py | 140 ++--- src/finch/tensor.py | 1143 +++++++++++++++++----------------- src/finch/typing.py | 8 +- tests/test_compiler.py | 192 ++++++ 10 files changed, 1430 insertions(+), 1173 deletions(-) create mode 100644 tests/test_compiler.py diff --git a/src/finch/__init__.py b/src/finch/__init__.py index c853a41..eff4ae6 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -1,359 +1,362 @@ -from operator import ( - abs as abs, -) -from operator import ( - add as add, -) -from operator import ( - and_ as bitwise_and, -) -from operator import ( - eq as equal, -) -from operator import ( - floordiv as floor_divide, -) -from operator import ( - ge as greater_equal, -) -from operator import ( - gt as greater, -) -from operator import ( - invert as bitwise_invert, -) -from operator import ( - le as less_equal, -) -from operator import ( - lshift as bitwise_left_shift, -) -from operator import ( - lt as less, -) -from operator import ( - matmul as matmul, -) -from operator import ( - mod as remainder, -) -from operator import ( - mul as multiply, -) -from operator import ( - ne as not_equal, -) -from operator import ( - neg as negative, -) -from operator import ( - or_ as bitwise_or, -) -from operator import ( - pos as positive, -) -from operator import ( - pow as pow, -) -from operator import ( - rshift as bitwise_right_shift, -) -from operator import ( - sub as subtract, -) -from operator import ( - truediv as divide, -) -from operator import ( - xor as bitwise_xor, -) +# from operator import ( +# abs as abs, +# ) +# from operator import ( +# add as add, +# ) +# from operator import ( +# and_ as bitwise_and, +# ) +# from operator import ( +# eq as equal, +# ) +# from operator import ( +# floordiv as floor_divide, +# ) +# from operator import ( +# ge as greater_equal, +# ) +# from operator import ( +# gt as greater, +# ) +# from operator import ( +# invert as bitwise_invert, +# ) +# from operator import ( +# le as less_equal, +# ) +# from operator import ( +# lshift as bitwise_left_shift, +# ) +# from operator import ( +# lt as less, +# ) +# from operator import ( +# matmul as matmul, +# ) +# from operator import ( +# mod as remainder, +# ) +# from operator import ( +# mul as multiply, +# ) +# from operator import ( +# ne as not_equal, +# ) +# from operator import ( +# neg as negative, +# ) +# from operator import ( +# or_ as bitwise_or, +# ) +# from operator import ( +# pos as positive, +# ) +# from operator import ( +# pow as pow, +# ) +# from operator import ( +# rshift as bitwise_right_shift, +# ) +# from operator import ( +# sub as subtract, +# ) +# from operator import ( +# truediv as divide, +# ) +# from operator import ( +# xor as bitwise_xor, +# ) -from numpy import ( - e as e, -) -from numpy import ( - inf as inf, -) -from numpy import ( - nan as nan, -) -from numpy import ( - newaxis as newaxis, -) -from numpy import ( - pi as pi, -) +# from numpy import ( +# e as e, +# ) +# from numpy import ( +# inf as inf, +# ) +# from numpy import ( +# nan as nan, +# ) +# from numpy import ( +# newaxis as newaxis, +# ) +# from numpy import ( +# pi as pi, +# ) -from . import linalg -from ._array_api_info import __array_namespace_info__ -from .compiled import ( - DefaultScheduler, - GalleyScheduler, - compiled, - compute, - lazy, - set_optimizer, -) -from .dtypes import ( - bool, - can_cast, - complex64, - complex128, - finfo, - float16, - float32, - float64, - iinfo, - int8, - int16, - int32, - int64, - int_, - uint, - uint8, - uint16, - uint32, - uint64, -) -from .io import ( - read, - write, -) -from .levels import ( - Dense, - DenseStorage, - Element, - Pattern, - RepeatRLE, - SparseByteMap, - SparseCOO, - SparseHash, - SparseList, - SparseVBL, - Storage, -) -from .tensor import ( - SparseArray, - FinchJLTensor, - acos, - acosh, - all, - any, - arange, - argmax, - argmin, - asarray, - asin, - asinh, - astype, - atan, - atan2, - atanh, - ceil, - conj, - cos, - cosh, - diagonal, - einop, - einsum, - empty, - empty_like, - exp, - expand_dims, - expm1, - eye, - floor, - full, - full_like, - imag, - isfinite, - isinf, - isnan, - linspace, - log, - log1p, - log2, - log10, - logaddexp, - logical_and, - logical_or, - logical_xor, - max, - mean, - min, - moveaxis, - nonzero, - ones, - ones_like, - permute_dims, - power, - prod, - random, - real, - reshape, - round, - sign, - sin, - sinh, - sqrt, - square, - squeeze, - std, - sum, - tan, - tanh, - tensordot, - trunc, - var, - where, - zeros, - zeros_like, -) +# from . import linalg +# from ._array_api_info import __array_namespace_info__ +# from .compiled import ( +# DefaultScheduler, +# GalleyScheduler, +# compiled, +# compute, +# lazy, +# set_optimizer, +# ) +# from .dtypes import ( +# bool, +# can_cast, +# complex64, +# complex128, +# finfo, +# float16, +# float32, +# float64, +# iinfo, +# int8, +# int16, +# int32, +# int64, +# int_, +# uint, +# uint8, +# uint16, +# uint32, +# uint64, +# ) +# from .io import ( +# read, +# write, +# ) +# from .levels import ( +# Dense, +# DenseStorage, +# Element, +# Pattern, +# RepeatRLE, +# SparseByteMap, +# SparseCOO, +# SparseHash, +# SparseList, +# SparseVBL, +# Storage, +# ) +# from .tensor import ( +# SparseArray, +# FinchJLTensor, +# acos, +# acosh, +# all, +# any, +# arange, +# argmax, +# argmin, +# asarray, +# asin, +# asinh, +# astype, +# atan, +# atan2, +# atanh, +# ceil, +# conj, +# cos, +# cosh, +# diagonal, +# einop, +# einsum, +# empty, +# empty_like, +# exp, +# expand_dims, +# expm1, +# eye, +# floor, +# full, +# full_like, +# imag, +# isfinite, +# isinf, +# isnan, +# linspace, +# log, +# log1p, +# log2, +# log10, +# logaddexp, +# logical_and, +# logical_or, +# logical_xor, +# max, +# mean, +# min, +# moveaxis, +# nonzero, +# ones, +# ones_like, +# permute_dims, +# power, +# prod, +# random, +# real, +# reshape, +# round, +# sign, +# sin, +# sinh, +# sqrt, +# square, +# squeeze, +# std, +# sum, +# tan, +# tanh, +# tensordot, +# trunc, +# var, +# where, +# zeros, +# zeros_like, +# ) -__all__ = [ - "DefaultScheduler", - "Dense", - "DenseStorage", - "Element", - "GalleyScheduler", - "Pattern", - "RepeatRLE", - "SparseArray", - "SparseByteMap", - "SparseCOO", - "SparseHash", - "SparseList", - "SparseVBL", - "Storage", - "FinchJLTensor", - "__array_namespace_info__", - "abs", - "acos", - "acosh", - "add", - "all", - "any", - "arange", - "argmax", - "argmin", - "asarray", - "asin", - "asinh", - "astype", - "atan", - "atan2", - "atanh", - "bitwise_and", - "bitwise_invert", - "bitwise_left_shift", - "bitwise_or", - "bitwise_right_shift", - "bitwise_xor", - "bool", - "can_cast", - "ceil", - "compiled", - "complex64", - "complex128", - "compute", - "conj", - "cos", - "cosh", - "diagonal", - "divide", - "e", - "einop", - "einsum", - "empty", - "empty_like", - "equal", - "exp", - "expand_dims", - "expm1", - "eye", - "finfo", - "float16", - "float32", - "float64", - "floor", - "floor_divide", - "full", - "full_like", - "greater", - "greater_equal", - "iinfo", - "imag", - "inf", - "int8", - "int16", - "int32", - "int64", - "int_", - "isfinite", - "isinf", - "isnan", - "lazy", - "less", - "less_equal", - "linalg", - "linspace", - "log", - "log1p", - "log2", - "log10", - "logaddexp", - "logical_and", - "logical_or", - "logical_xor", - "matmul", - "max", - "mean", - "min", - "moveaxis", - "multiply", - "nan", - "negative", - "newaxis", - "nonzero", - "not_equal", - "ones", - "ones_like", - "permute_dims", - "pi", - "positive", - "pow", - "power", - "prod", - "random", - "read", - "real", - "remainder", - "reshape", - "round", - "set_optimizer", - "sign", - "sin", - "sinh", - "sqrt", - "square", - "squeeze", - "std", - "subtract", - "sum", - "tan", - "tanh", - "tensordot", - "trunc", - "uint", - "uint8", - "uint16", - "uint32", - "uint64", - "var", - "where", - "write", - "zeros", - "zeros_like", -] +# __all__ = [ +# "DefaultScheduler", +# "Dense", +# "DenseStorage", +# "Element", +# "GalleyScheduler", +# "Pattern", +# "RepeatRLE", +# "SparseArray", +# "SparseByteMap", +# "SparseCOO", +# "SparseHash", +# "SparseList", +# "SparseVBL", +# "Storage", +# "FinchJLTensor", +# "__array_namespace_info__", +# "abs", +# "acos", +# "acosh", +# "add", +# "all", +# "any", +# "arange", +# "argmax", +# "argmin", +# "asarray", +# "asin", +# "asinh", +# "astype", +# "atan", +# "atan2", +# "atanh", +# "bitwise_and", +# "bitwise_invert", +# "bitwise_left_shift", +# "bitwise_or", +# "bitwise_right_shift", +# "bitwise_xor", +# "bool", +# "can_cast", +# "ceil", +# "compiled", +# "complex64", +# "complex128", +# "compute", +# "conj", +# "cos", +# "cosh", +# "diagonal", +# "divide", +# "e", +# "einop", +# "einsum", +# "empty", +# "empty_like", +# "equal", +# "exp", +# "expand_dims", +# "expm1", +# "eye", +# "finfo", +# "float16", +# "float32", +# "float64", +# "floor", +# "floor_divide", +# "full", +# "full_like", +# "greater", +# "greater_equal", +# "iinfo", +# "imag", +# "inf", +# "int8", +# "int16", +# "int32", +# "int64", +# "int_", +# "isfinite", +# "isinf", +# "isnan", +# "lazy", +# "less", +# "less_equal", +# "linalg", +# "linspace", +# "log", +# "log1p", +# "log2", +# "log10", +# "logaddexp", +# "logical_and", +# "logical_or", +# "logical_xor", +# "matmul", +# "max", +# "mean", +# "min", +# "moveaxis", +# "multiply", +# "nan", +# "negative", +# "newaxis", +# "nonzero", +# "not_equal", +# "ones", +# "ones_like", +# "permute_dims", +# "pi", +# "positive", +# "pow", +# "power", +# "prod", +# "random", +# "read", +# "real", +# "remainder", +# "reshape", +# "round", +# "set_optimizer", +# "sign", +# "sin", +# "sinh", +# "sqrt", +# "square", +# "squeeze", +# "std", +# "subtract", +# "sum", +# "tan", +# "tanh", +# "tensordot", +# "trunc", +# "uint", +# "uint8", +# "uint16", +# "uint32", +# "uint64", +# "var", +# "where", +# "write", +# "zeros", +# "zeros_like", +# ] -__array_api_version__: str = "2024.12" +# __array_api_version__: str = "2024.12" + + +from .compiler import FinchJLCompiler diff --git a/src/finch/_array_api_info.py b/src/finch/_array_api_info.py index c4e3f5a..b8d57c0 100644 --- a/src/finch/_array_api_info.py +++ b/src/finch/_array_api_info.py @@ -1,94 +1,94 @@ -from . import dtypes -from .typing import DType +# from . import dtypes +# from .typing import DType -class __array_namespace_info__: - def capabilities(self) -> dict[str, bool]: - return { - "boolean indexing": True, - "data-dependent shapes": True, - } +# class __array_namespace_info__: +# def capabilities(self) -> dict[str, bool]: +# return { +# "boolean indexing": True, +# "data-dependent shapes": True, +# } - def default_device(self) -> str: - return "cpu" +# def default_device(self) -> str: +# return "cpu" - def default_dtypes(self, *, device: str | None = None) -> dict[str, DType]: - if device not in ["cpu", None]: - raise ValueError( - f'Device not understood. Only "cpu" is allowed, but received: {device}' - ) - return { - "real floating": dtypes.float64, - "complex floating": dtypes.complex128, - "integral": dtypes.int_, - "indexing": dtypes.int_, - } +# def default_dtypes(self, *, device: str | None = None) -> dict[str, DType]: +# if device not in ["cpu", None]: +# raise ValueError( +# f'Device not understood. Only "cpu" is allowed, but received: {device}' +# ) +# return { +# "real floating": dtypes.float64, +# "complex floating": dtypes.complex128, +# "integral": dtypes.int_, +# "indexing": dtypes.int_, +# } - _bool_dtypes = {"bool": dtypes.bool} - _signed_integer_dtypes = { - "int8": dtypes.int8, - "int16": dtypes.int16, - "int32": dtypes.int32, - "int64": dtypes.int64, - } - _unsigned_integer_dtypes = { - "uint8": dtypes.uint8, - "uint16": dtypes.uint16, - "uint32": dtypes.uint32, - "uint64": dtypes.uint64, - } - _real_floating_dtypes = { - "float32": dtypes.float32, - "float64": dtypes.float64, - } - _complex_floating_dtypes = { - "complex64": dtypes.complex64, - "complex128": dtypes.complex128, - } +# _bool_dtypes = {"bool": dtypes.bool} +# _signed_integer_dtypes = { +# "int8": dtypes.int8, +# "int16": dtypes.int16, +# "int32": dtypes.int32, +# "int64": dtypes.int64, +# } +# _unsigned_integer_dtypes = { +# "uint8": dtypes.uint8, +# "uint16": dtypes.uint16, +# "uint32": dtypes.uint32, +# "uint64": dtypes.uint64, +# } +# _real_floating_dtypes = { +# "float32": dtypes.float32, +# "float64": dtypes.float64, +# } +# _complex_floating_dtypes = { +# "complex64": dtypes.complex64, +# "complex128": dtypes.complex128, +# } - def dtypes( - self, - *, - device: str | None = None, - kind: str | tuple[str, ...] | None = None, - ) -> dict[str, DType]: - if device not in ["cpu", None]: - raise ValueError( - f'Device not understood. Only "cpu" is allowed, but received: {device}' - ) - if kind is None: - return ( - self._bool_dtypes - | self._signed_integer_dtypes - | self._unsigned_integer_dtypes - | self._real_floating_dtypes - | self._complex_floating_dtypes - ) - if kind == "bool": - return self._bool_dtypes - if kind == "signed integer": - return self._signed_integer_dtypes - if kind == "unsigned integer": - return self._unsigned_integer_dtypes - if kind == "integral": - return self._signed_integer_dtypes | self._unsigned_integer_dtypes - if kind == "real floating": - return self._real_floating_dtypes - if kind == "complex floating": - return self._complex_floating_dtypes - if kind == "numeric": - return ( - self._signed_integer_dtypes - | self._unsigned_integer_dtypes - | self._real_floating_dtypes - | self._complex_floating_dtypes - ) - if isinstance(kind, tuple): - res = {} - for k in kind: - res.update(self.dtypes(kind=k)) - return res - raise ValueError(f"unsupported kind: {kind!r}") +# def dtypes( +# self, +# *, +# device: str | None = None, +# kind: str | tuple[str, ...] | None = None, +# ) -> dict[str, DType]: +# if device not in ["cpu", None]: +# raise ValueError( +# f'Device not understood. Only "cpu" is allowed, but received: {device}' +# ) +# if kind is None: +# return ( +# self._bool_dtypes +# | self._signed_integer_dtypes +# | self._unsigned_integer_dtypes +# | self._real_floating_dtypes +# | self._complex_floating_dtypes +# ) +# if kind == "bool": +# return self._bool_dtypes +# if kind == "signed integer": +# return self._signed_integer_dtypes +# if kind == "unsigned integer": +# return self._unsigned_integer_dtypes +# if kind == "integral": +# return self._signed_integer_dtypes | self._unsigned_integer_dtypes +# if kind == "real floating": +# return self._real_floating_dtypes +# if kind == "complex floating": +# return self._complex_floating_dtypes +# if kind == "numeric": +# return ( +# self._signed_integer_dtypes +# | self._unsigned_integer_dtypes +# | self._real_floating_dtypes +# | self._complex_floating_dtypes +# ) +# if isinstance(kind, tuple): +# res = {} +# for k in kind: +# res.update(self.dtypes(kind=k)) +# return res +# raise ValueError(f"unsupported kind: {kind!r}") - def devices(self) -> list[str]: - return ["cpu"] +# def devices(self) -> list[str]: +# return ["cpu"] diff --git a/src/finch/compiler.py b/src/finch/compiler.py index d0bbee6..10281cc 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,29 +1,30 @@ -from .tensor import FinchJLTensor -from finchlite import NotationCompiler, AssemblyKernel, AssemblyLibrary +from finchlite.compile import NotationCompiler +from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary import finchlite.finch_notation.nodes as ntn +from typing import Any -from juliacall import Main as jl +from juliacall import Main as jl -CompiledModuleName = "compiled_module" -FunctionName = "compiled_function" class FinchJLKernel(AssemblyKernel): - def __init__(self, func_name, jl_code): + def __init__(self, func_name, jl_code): + self.jl_code = jl_code self.func_name = func_name jl.seval(jl_code) - def __call__( - self, *args: tuple[FinchJLTensor, ...] - ) -> tuple[FinchJLTensor, ...]: - argList = [] + + # TODO: Switch back to (self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...] + def __call__(self, *args: tuple[Any, ...]): + argList = [] for arg in args: argList.append(f"arg{len(argList)}") setattr(jl, argList[-1], arg) jl.seval(f"{self.func_name}({argList.join(',')})") + class FinchJLLibrary(AssemblyLibrary): - def __init__(self, kernel_name, kernel): - self.kernel_dict = {kernel_name: kernel} + def __init__(self, kernel_dict): + self.kernel_dict = kernel_dict def __getattr__(self, name: str) -> FinchJLKernel: return self.kernels[name] @@ -32,11 +33,67 @@ def __getattr__(self, name: str) -> FinchJLKernel: class FinchJLGenerator: def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: match prgm: - case ntn.Module(statements=stmts): + case ntn.Literal(val): + ... + # What this? + case ntn.Value(ex, type_): + ... + case ntn.Variable(name, type_): + ... + case ntn.Call(op, args): + ... + case ntn.Dimension(tns, r): + ... + case ntn.Access(tns, mode, idxs): + ... + case ntn.Read(): + ... + case ntn.AccessMode(): + ... + case ntn.Update(op): + ... + case ntn.Increment(lhs,rhs): + ... + case ntn.Unwrap(arg): + ... + case ntn.Cached(arg,ref): + ... + case ntn.Loop(idx, ext, body): + ... + case ntn.If(cond, body): + ... + case ntn.IfElse(cond, then_body, else_body): + ... + case ntn.Assign(lhs,rhs): + ... + case ntn.Stack(obj,type): + ... + case ntn.Slot(name, type): + ... + case ntn.Unpack(lhs, rhs): + ... + case ntn.Repack(val, obj): + ... + case ntn.Declare(tns, init, op, shape): + ... + case ntn.Freeze(tns, op): + ... + case ntn.Thaw(tns, op): + ... + case ntn.Block(bodies): + ... + case ntn.Function(name,args,body): + ... + case ntn.Return(val): ... + class FinchJLCompiler(NotationCompiler): def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: generator = FinchJLGenerator() - jl_code = generator(prgm) - return FinchJLLibrary(CompiledModuleName, FinchJLKernel(FunctionName, jl_code)) \ No newline at end of file + + kernel_dict = {} + for func in prgm.children: + kernel_dict[func.name.name] = FinchJLKernel(func.name.name, generator(func)) + + return FinchJLLibrary(kernel_dict) diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py index 2b95252..3b6b4cc 100644 --- a/src/finch/dtypes.py +++ b/src/finch/dtypes.py @@ -1,61 +1,61 @@ -import builtins - -import numpy as np - -from .julia import jl - -int_: jl.DataType = jl.Int -int8: jl.DataType = jl.Int8 -int16: jl.DataType = jl.Int16 -int32: jl.DataType = jl.Int32 -int64: jl.DataType = jl.Int64 -uint: jl.DataType = jl.UInt -uint8: jl.DataType = jl.UInt8 -uint16: jl.DataType = jl.UInt16 -uint32: jl.DataType = jl.UInt32 -uint64: jl.DataType = jl.UInt64 -float16: jl.DataType = jl.Float16 -float32: jl.DataType = jl.Float32 -float64: jl.DataType = jl.Float64 -complex64: jl.DataType = jl.ComplexF32 -complex128: jl.DataType = jl.ComplexF64 -bool: jl.DataType = jl.Bool - -number: jl.DataType = jl.Number -complex: jl.DataType = jl.Complex -integer: jl.DataType = jl.Integer -abstract_float: jl.DataType = jl.AbstractFloat - -jl_to_np_dtype = { - int_: np.int_, - int8: np.int8, - int16: np.int16, - int32: np.int32, - int64: np.int64, - uint: np.uint, - uint8: np.uint8, - uint16: np.uint16, - uint32: np.uint32, - uint64: np.uint64, - float16: np.float16, - float32: np.float32, - float64: np.float64, - complex64: np.complex64, - complex128: np.complex128, - bool: builtins.bool, - None: None, -} - - -def finfo(dtype): - return np.finfo(jl_to_np_dtype[dtype]) - - -def iinfo(dtype): - return np.iinfo(jl_to_np_dtype[dtype]) - - -def can_cast(from_, to, /) -> builtins.bool: - if hasattr(from_, "dtype"): - from_ = from_.dtype - return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) +# import builtins + +# import numpy as np + +# from .julia import jl + +# int_: jl.DataType = jl.Int +# int8: jl.DataType = jl.Int8 +# int16: jl.DataType = jl.Int16 +# int32: jl.DataType = jl.Int32 +# int64: jl.DataType = jl.Int64 +# uint: jl.DataType = jl.UInt +# uint8: jl.DataType = jl.UInt8 +# uint16: jl.DataType = jl.UInt16 +# uint32: jl.DataType = jl.UInt32 +# uint64: jl.DataType = jl.UInt64 +# float16: jl.DataType = jl.Float16 +# float32: jl.DataType = jl.Float32 +# float64: jl.DataType = jl.Float64 +# complex64: jl.DataType = jl.ComplexF32 +# complex128: jl.DataType = jl.ComplexF64 +# bool: jl.DataType = jl.Bool + +# number: jl.DataType = jl.Number +# complex: jl.DataType = jl.Complex +# integer: jl.DataType = jl.Integer +# abstract_float: jl.DataType = jl.AbstractFloat + +# jl_to_np_dtype = { +# int_: np.int_, +# int8: np.int8, +# int16: np.int16, +# int32: np.int32, +# int64: np.int64, +# uint: np.uint, +# uint8: np.uint8, +# uint16: np.uint16, +# uint32: np.uint32, +# uint64: np.uint64, +# float16: np.float16, +# float32: np.float32, +# float64: np.float64, +# complex64: np.complex64, +# complex128: np.complex128, +# bool: builtins.bool, +# None: None, +# } + + +# def finfo(dtype): +# return np.finfo(jl_to_np_dtype[dtype]) + + +# def iinfo(dtype): +# return np.iinfo(jl_to_np_dtype[dtype]) + + +# def can_cast(from_, to, /) -> builtins.bool: +# if hasattr(from_, "dtype"): +# from_ = from_.dtype +# return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) diff --git a/src/finch/errors.py b/src/finch/errors.py index c034e04..3a79630 100644 --- a/src/finch/errors.py +++ b/src/finch/errors.py @@ -1,2 +1,2 @@ -class PerformanceWarning(Warning): - pass +# class PerformanceWarning(Warning): +# pass diff --git a/src/finch/io.py b/src/finch/io.py index 700e2d6..4e10a72 100644 --- a/src/finch/io.py +++ b/src/finch/io.py @@ -1,15 +1,15 @@ -from pathlib import Path +# from pathlib import Path -from .julia import jl -from .tensor import Tensor +# from .julia import jl +# from .tensor import Tensor -def read(filename: Path | str) -> Tensor: - fn = str(filename) - julia_obj = jl.fread(fn) - return Tensor(julia_obj) +# def read(filename: Path | str) -> Tensor: +# fn = str(filename) +# julia_obj = jl.fread(fn) +# return Tensor(julia_obj) -def write(filename: Path | str, tns: Tensor) -> None: - fn = str(filename) - jl.fwrite(fn, tns._obj) +# def write(filename: Path | str, tns: Tensor) -> None: +# fn = str(filename) +# jl.fwrite(fn, tns._obj) diff --git a/src/finch/levels.py b/src/finch/levels.py index 07d6142..f65a5ad 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -1,110 +1,110 @@ -from .julia import jl -from .typing import DType, JuliaObj, OrderType +# from .julia import jl +# from .typing import DType, JuliaObj, OrderType -class _Display: - _obj: JuliaObj +# class _Display: +# _obj: JuliaObj - def __repr__(self): - return jl.sprint(jl.show, self._obj) +# def __repr__(self): +# return jl.sprint(jl.show, self._obj) - def __str__(self): - return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) +# def __str__(self): +# return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) -# LEVEL +# # LEVEL -class AbstractLevel(_Display): - pass +# class AbstractLevel(_Display): +# pass -# core levels +# # core levels -class Dense(AbstractLevel): - def __init__(self, lvl, shape=None): - args = [lvl._obj] - if shape is not None: - args.append(shape) - self._obj = jl.Dense(*args) +# class Dense(AbstractLevel): +# def __init__(self, lvl, shape=None): +# args = [lvl._obj] +# if shape is not None: +# args.append(shape) +# self._obj = jl.Dense(*args) -class Element(AbstractLevel): - def __init__(self, fill_value, data=None): - args = [fill_value] - if data is not None: - args.append(data) - self._obj = jl.Element(*args) +# class Element(AbstractLevel): +# def __init__(self, fill_value, data=None): +# args = [fill_value] +# if data is not None: +# args.append(data) +# self._obj = jl.Element(*args) -class Pattern(AbstractLevel): - def __init__(self): - self._obj = jl.Pattern() +# class Pattern(AbstractLevel): +# def __init__(self): +# self._obj = jl.Pattern() -# advanced levels +# # advanced levels -class SparseList(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseList(lvl._obj) +# class SparseList(AbstractLevel): +# def __init__(self, lvl): +# self._obj = jl.SparseList(lvl._obj) -class SparseByteMap(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseByteMap(lvl._obj) +# class SparseByteMap(AbstractLevel): +# def __init__(self, lvl): +# self._obj = jl.SparseByteMap(lvl._obj) -class RepeatRLE(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.RepeatRLE(lvl._obj) +# class RepeatRLE(AbstractLevel): +# def __init__(self, lvl): +# self._obj = jl.RepeatRLE(lvl._obj) -class SparseVBL(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseVBL(lvl._obj) +# class SparseVBL(AbstractLevel): +# def __init__(self, lvl): +# self._obj = jl.SparseVBL(lvl._obj) -class SparseCOO(AbstractLevel): - def __init__(self, ndim, lvl): - self._obj = jl.SparseCOO[ndim](lvl._obj) +# class SparseCOO(AbstractLevel): +# def __init__(self, ndim, lvl): +# self._obj = jl.SparseCOO[ndim](lvl._obj) -class SparseHash(AbstractLevel): - def __init__(self, ndim, lvl): - self._obj = jl.SparseHash[ndim](lvl._obj) +# class SparseHash(AbstractLevel): +# def __init__(self, ndim, lvl): +# self._obj = jl.SparseHash[ndim](lvl._obj) -sparse_formats_names = ( - "SparseList", - "Sparse", - "SparseHash", - "SparseCOO", - "SparseRLE", - "SparseVBL", - "SparseBand", - "SparsePoint", - "SparseInterval", -) +# sparse_formats_names = ( +# "SparseList", +# "Sparse", +# "SparseHash", +# "SparseCOO", +# "SparseRLE", +# "SparseVBL", +# "SparseBand", +# "SparsePoint", +# "SparseInterval", +# ) -# STORAGE +# # STORAGE -class Storage: - def __init__(self, levels_descr: AbstractLevel, order: OrderType = None): - self.levels_descr = levels_descr - self.order = order if order is not None else "C" +# class Storage: +# def __init__(self, levels_descr: AbstractLevel, order: OrderType = None): +# self.levels_descr = levels_descr +# self.order = order if order is not None else "C" - def __str__(self) -> str: - return f"Storage(lvl={str(self.levels_descr)}, order={self.order})" +# def __str__(self) -> str: +# return f"Storage(lvl={str(self.levels_descr)}, order={self.order})" -class DenseStorage(Storage): - def __init__(self, ndim: int, dtype: DType, order: OrderType = None): - lvl = Element(dtype(0)) - for _ in range(ndim): - lvl = Dense(lvl) +# class DenseStorage(Storage): +# def __init__(self, ndim: int, dtype: DType, order: OrderType = None): +# lvl = Element(dtype(0)) +# for _ in range(ndim): +# lvl = Dense(lvl) - super().__init__(levels_descr=lvl, order=order) +# super().__init__(levels_descr=lvl, order=order) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 5bcf328..afdffb1 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,570 +1,577 @@ -from __future__ import annotations - -import builtins -import warnings -from collections.abc import Callable, Iterable -from typing import Any, Literal - -import numpy as np -from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple - -from . import dtypes as jl_dtypes -from .errors import PerformanceWarning -from .julia import jc, jl -from .levels import ( - Dense, - DenseStorage, - Element, - SparseCOO, - SparseList, - Storage, - _Display, - sparse_formats_names, -) -from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix -from finchlite import Tensor, TensorFType - - -class SparseArray: - """ - PyData/Sparse marker class - """ - -class FinchJLTensorFType(TensorFType): - def __init__(self, jltype): - self.jltype = jltype - - def ndims(self) -> int: - return jl.ndims(self.jltype) +# from __future__ import annotations + +# import builtins +# import warnings +# from collections.abc import Callable, Iterable +# from typing import Any, Literal + +# import numpy as np +# from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple + +# from . import dtypes as jl_dtypes +# from .errors import PerformanceWarning +# from .julia import jc, jl +# from .levels import ( +# Dense, +# DenseStorage, +# Element, +# SparseCOO, +# SparseList, +# Storage, +# _Display, +# sparse_formats_names, +# ) +# from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix +# from finchlite import Tensor, TensorFType, EagerTensor + + +# class FinchJLTensorFType(TensorFType): +# def __init__(self, jltype): +# # Julia type associated with the tensor +# self.jltype = jltype + +# def ndims(self) -> np.intp: +# return np.intp(jl.ndims(self.jltype)) - def element_type(self): - return jl.eltype(self.jltype) +# def fill_value(self) -> Any: +# return jl.fill_value(self.jltype) - ... - -class FinchJLTensor(_Display, SparseArray, Tensor, finchlite.EagerTensor): - """ - A wrapper class for Finch.Tensor and Finch.SwizzleArray. - - Constructors - ------------ - FinchJLTensor(scipy.sparse.spmatrix) - Construct a Tensor out of a `scipy.sparse` object. Supported formats are: `COO`, - `CSC`, and `CSR`. - FinchJLTensor(numpy.ndarray) - Construct a Tensor out of a NumPy array object. This is a no-copy operation. - FinchJLTensor(Storage) - Initialize a Tensor with a `storage` description. `storage` can already hold - data. - FinchJLTensor(julia_object) - Tensor created from a compatible raw Julia object. Must be a `Tensor`. - This is a no-copy operation. - - Parameters - ---------- - obj : np.ndarray or scipy.sparse or Storage or Finch.Tensor - Input to construct a Tensor. It's a no-copy operation of for NumPy and - SciPy input. For Storage it's levels' description with order. The order - numbers the dimensions from the fastest to slowest. The leaf nodes have - mode `0` and the root node has mode `n-1`. If the tensor was square of - size `N`, then `N .^ order == strides`. Available options are "C" - (row-major), "F" (column-major), or a custom order. Default: row-major. - fill_value : np.number, optional - Only used when `numpy.ndarray` or `scipy.sparse` is passed. - copy : bool, optional - If ``True``, then the object is copied. If ``None`` then the object is - copied only if needed. For ``False`` it raises a ``ValueError`` if a - copy cannot be avoided. Default: ``None``. - - Returns - ------- - FinchJLTensor - Python wrapper for Finch.jl `Tensor`. - - Examples - -------- - >>> import numpy as np - >>> import finch - >>> arr2d = np.arange(6).reshape((2, 3)) - >>> t1 = finch.FinchJLTensor(arr2d) - >>> t1.todense() - array([[0, 1, 2], - [3, 4, 5]]) - >>> np.shares_memory(t1.todense(), arr2d) - True - >>> storage = finch.Storage( - ... finch.Dense(finch.SparseList(finch.Element(1))), order="C" - ... ) - >>> t2 = t1.to_storage(storage) - >>> t2.todense() - array([[0, 1, 2], - [3, 4, 5]]) - """ - - def __init__( - self, - obj: np.ndarray | spmatrix | Storage | JuliaObj, - /, - *, - fill_value: np.number | None = None, - copy: bool | None = None, - ): - if isinstance(obj, int | float | complex | bool | list): - if copy is False: - raise ValueError( - "copy=False isn't supported for scalar inputs and Python lists" - ) - obj = np.asarray(obj) - if fill_value is None: - fill_value = 0.0 - - if _is_scipy_sparse_obj(obj): # scipy constructor - jl_data = self._from_scipy_sparse(obj, fill_value=fill_value, copy=copy) - self._obj = jl_data - elif isinstance(obj, np.ndarray): # numpy constructor - jl_data = self._from_numpy(obj, fill_value=fill_value, copy=copy) - self._obj = jl_data - elif isinstance(obj, Storage): # from-storage constructor - if copy: - self._raise_julia_copy_not_supported() - order = self.preprocess_order( - obj.order, self.get_lvl_ndim(obj.levels_descr._obj) - ) - self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order) - elif jl.isa(obj, jl.Finch.Tensor): # raw-Julia-object constructors - if copy: - self._raise_julia_copy_not_supported() - self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1))) - elif jl.isa(obj, jl.Finch.Tensor): - if copy: - self._raise_julia_copy_not_supported() - self._obj = obj - elif isinstance(obj, FinchJLTensor): - self._obj = obj._obj - else: - raise ValueError( - "Either scalar, numpy, scipy.sparse or a raw julia object should " - f"be provided. Found: {type(obj)}" - ) - - @property - def element_type(self): - return jl.eltype(self._obj.body) - - @property - def dtype(self) -> DType: - return jl.eltype(self._obj.body) - - @property - def ndim(self) -> int: - return jl.ndims(self._obj) - - @property - def shape(self) -> tuple[int, ...]: - return jl.size(self._obj) - - @property - def size(self) -> int: - return np.prod(self.shape) - - @property - def fill_value(self) -> np.number: - return jl.fill_value(self._obj) - - @property - def _is_dense(self) -> bool: - lvl = self._obj.body.lvl - for _ in self.shape: - if not jl.isa(lvl, jl.Finch.Dense): - return False - lvl = lvl.lvl - return True - - @property - def _order(self) -> tuple[int, ...]: - return jl.typeof(self._obj).parameters[1] - - @property - def mT(self) -> Tensor: - axes = list(range(self.ndim)) - axes[-2], axes[-1] = axes[-1], axes[-2] - axes = tuple(axes) - return self.permute_dims(axes) - - @property - def device(self) -> str: - return "cpu" - - def to_device( - self, device: Device, /, *, stream: int | Any | None = None - ) -> Tensor: - if device != "cpu": - raise ValueError("Only `device='cpu'` is supported.") - - return self - - @classmethod - def get_lvl_ndim(cls, lvl: JuliaObj) -> int: - ndim = 0 - while True: - ndim += 1 - lvl = lvl.lvl - if jl.isa(lvl, jl.Finch.Element): - break - return ndim - - def todense(self) -> np.ndarray: - obj = self._obj - - if self._is_dense: - # don't materialize a dense finch tensor - shape = jl.size(obj.body) - dense_tensor = obj.body.lvl - else: - # create materialized dense array - shape = jl.size(obj) - dense_lvls = jl.Element(jc.convert(self.dtype, jl.fill_value(obj))) - for _ in range(self.ndim): - dense_lvls = jl.Dense(dense_lvls) - dense_tensor = jl.Tensor(dense_lvls, obj).lvl # materialize - - for _ in range(self.ndim): - dense_tensor = dense_tensor.lvl - - result = np.asarray(jl.reshape(dense_tensor.val, shape)) - return result.transpose(self.get_order()) if self._is_dense else result - - #TODO: Do we need? - def permute_dims(self, axes: tuple[int, ...]) -> Tensor: - axes = tuple(i + 1 for i in axes) - new_obj = jl.permutedims(self._obj, axes) - return Tensor(new_obj) - - def to_storage(self, storage: Storage) -> Tensor: - return Tensor(self._from_other_tensor(self, storage=storage)) - - @classmethod - def _from_other_tensor(cls, tensor: Tensor, storage: Storage) -> JuliaObj: - order = cls.preprocess_order(storage.order, tensor.ndim) - result = jl.copyto_b( - jl.swizzle(jl.Tensor(storage.levels_descr._obj), *order), tensor._obj - ) - return jl.dropfills(result) if tensor._is_dense else result - - @classmethod - def _from_numpy( - cls, arr: np.ndarray, fill_value: np.number, copy: bool | None = None - ) -> JuliaObj: - if copy: - arr = arr.copy() - order_char = "F" if np.isfortran(arr) else "C" - order = cls.preprocess_order(order_char, arr.ndim) - inv_order = tuple(i - 1 for i in jl.invperm(order)) - - dtype = arr.dtype.type - if ( - dtype == np.bool_ - ): # Fails with: Finch currently only supports isbits defaults - dtype = jl_dtypes.bool - fill_value = dtype(fill_value) - lvl = Element(fill_value, arr.reshape(-1, order=order_char)) - for i in inv_order: - lvl = Dense(lvl, arr.shape[i]) - return jl.swizzle(jl.Tensor(lvl._obj), *order) - - @classmethod - def from_scipy_sparse( - cls, - x, - fill_value: np.number | None = None, - copy: bool | None = None, - ) -> Tensor: - if not _is_scipy_sparse_obj(x): - raise ValueError("{x} is not a SciPy sparse object.") - return Tensor(x, fill_value=fill_value, copy=copy) - - @classmethod - def _from_scipy_sparse( - cls, - x, - *, - fill_value: np.number | None = None, - copy: bool | None = None, - ) -> JuliaObj: - if copy is False and not ( - x.format in ("coo", "csr", "csc") and x.has_canonical_format - ): - raise ValueError( - "Unable to avoid copy while creating an array as requested." - ) - if x.format not in ("coo", "csr", "csc"): - x = x.asformat("coo") - if copy: - x = x.copy() - if not x.has_canonical_format: - x.sum_duplicates() - assert x.has_canonical_format - - if x.format == "coo": - return cls.construct_coo_jl_object( - coords=(x.col, x.row), - data=x.data, - shape=x.shape[::-1], - order=Tensor.row_major, - fill_value=fill_value, - ) - if x.format == "csc": - return cls.construct_csc_jl_object( - arg=(x.data, x.indices, x.indptr), - shape=x.shape, - fill_value=fill_value, - ) - if x.format == "csr": - return cls.construct_csr_jl_object( - arg=(x.data, x.indices, x.indptr), - shape=x.shape, - fill_value=fill_value, - ) - raise ValueError(f"Unsupported SciPy format: {type(x)}") - - @classmethod - def construct_coo_jl_object( - cls, coords, data, shape, order, fill_value=0.0 - ) -> JuliaObj: - assert len(coords) == 2 - ndim = len(shape) - order = cls.preprocess_order(order, ndim) - - lvl = jl.Element(data.dtype.type(fill_value), data) - ptr = jl.Vector[jl.Int]([1, len(data) + 1]) - tbl = tuple(jl.PlusOneVector(arr) for arr in coords) - - return jl.swizzle(jl.Tensor(jl.SparseCOO[ndim](lvl, shape, ptr, tbl)), *order) - - @classmethod - def construct_coo( - cls, coords, data, shape, order=row_major, fill_value=0.0 - ) -> Tensor: - return Tensor( - cls.construct_coo_jl_object(coords, data, shape, order, fill_value) - ) - - @staticmethod - def _construct_compressed2d_jl_object( - arg: TupleOf3Arrays, - shape: tuple[int, ...], - order: tuple[int, ...], - fill_value: np.number = 0.0, - ) -> JuliaObj: - assert isinstance(arg, tuple) and len(arg) == 3 - assert len(shape) == 2 - - data, indices, indptr = arg - dtype = data.dtype.type - indices = jl.PlusOneVector(indices) - indptr = jl.PlusOneVector(indptr) - - lvl = jl.Element(dtype(fill_value), data) - return jl.swizzle( - jl.Tensor( - jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1]) - ), - *order, - ) - - @classmethod - def construct_csc_jl_object( - cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 - ) -> JuliaObj: - return cls._construct_compressed2d_jl_object( - arg=arg, shape=shape, order=(1, 2), fill_value=fill_value - ) - - @classmethod - def construct_csc( - cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 - ) -> Tensor: - return Tensor(cls.construct_csc_jl_object(arg, shape, fill_value)) - - @classmethod - def construct_csr_jl_object( - cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 - ) -> JuliaObj: - return cls._construct_compressed2d_jl_object( - arg=arg, shape=shape[::-1], order=(2, 1), fill_value=fill_value - ) - - @classmethod - def construct_csr( - cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 - ) -> Tensor: - return Tensor(cls.construct_csr_jl_object(arg, shape, fill_value)) - - @staticmethod - def construct_csf_jl_object( - arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 - ) -> JuliaObj: - assert isinstance(arg, tuple) and len(arg) == 3 - - data, indices_list, indptr_list = arg - dtype = data.dtype.type - - assert len(indices_list) == len(shape) - 1 - assert len(indptr_list) == len(shape) - 1 - - indices_list = [jl.PlusOneVector(i) for i in indices_list] - indptr_list = [jl.PlusOneVector(i) for i in indptr_list] - - lvl = jl.Element(dtype(fill_value), data) - for size, indices, indptr in zip( - shape[:-1], indices_list, indptr_list, strict=False - ): - lvl = jl.SparseList(lvl, size, indptr, indices) - - return jl.swizzle( - jl.Tensor(jl.Dense(lvl, shape[-1])), *range(1, len(shape) + 1) - ) - - @classmethod - def construct_csf( - cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 - ) -> Tensor: - return Tensor(cls.construct_csf_jl_object(arg, shape, fill_value)) - - def to_scipy_sparse(self, accept_fv=None): - import scipy.sparse as sp - - if accept_fv is None: - accept_fv = [0] - elif not isinstance(accept_fv, Iterable): - accept_fv = [accept_fv] - - if self.ndim != 2: - raise ValueError( - "Can only convert a 2-dimensional array to a Scipy sparse matrix." - ) - if not builtins.any(_eq_scalars(self.fill_value, fv) for fv in accept_fv): - raise ValueError( - f"Can only convert arrays with {accept_fv} fill-values " - "to a Scipy sparse matrix." - ) - order = self.get_order() - body = self._obj.body - - if str(jl.typeof(body.lvl).name.name) == "SparseCOOLevel": - data = np.asarray(body.lvl.lvl.val) - coords = body.lvl.tbl - row, col = coords[::-1] if order == (1, 0) else coords - row, col = np.asarray(row) - 1, np.asarray(col) - 1 - return sp.coo_matrix((data, (row, col)), shape=self.shape) - - if ( - str(jl.typeof(body.lvl).name.name) == "DenseLevel" - and str(jl.typeof(body.lvl.lvl).name.name) == "SparseListLevel" - ): - data = np.asarray(body.lvl.lvl.lvl.val) - indices = np.asarray(body.lvl.lvl.idx) - 1 - indptr = np.asarray(body.lvl.lvl.ptr) - 1 - sp_class = sp.csr_matrix if order == (1, 0) else sp.csc_matrix - return sp_class((data, indices, indptr), shape=self.shape) - if ( - jl.typeof(body.lvl).name.name in sparse_formats_names - or jl.typeof(body.lvl.lvl).name.name in sparse_formats_names - ): - storage = Storage(SparseCOO(self.ndim, Element(self.fill_value)), order) - return self.to_storage(storage).to_scipy_sparse() - raise ValueError("Tensor can't be converted to scipy.sparse object.") - - @staticmethod - def _raise_julia_copy_not_supported() -> None: - raise ValueError("copy=True isn't supported for Julia object inputs") - - def __array_namespace__(self, *, api_version: str | None = None) -> Any: - if api_version is None: - api_version = "2024.12" - - if api_version not in {"2021.12", "2022.12", "2023.12", "2024.12"}: - raise ValueError(f'"{api_version}" Array API version not supported.') - import finch - - return finch - - -def random(shape, density=0.01, random_state=None): - args = [*shape, density] - if random_state is not None: - if isinstance(random_state, np.random.Generator): - seed = random_state.integers(np.iinfo(np.int32).max) - else: - seed = random_state - rng = jl.Random.default_rng() - jl.Random.seed_b(rng, seed) - args = [rng] + args - return Tensor(jl.fsprand(*args)) - - -def asarray( - obj, - /, - *, - dtype: DType | None = None, - format: str | None = None, - fill_value: np.number | None = None, - device: Device | None = None, - copy: bool | None = None, -) -> Tensor: - if format not in {"coo", "csr", "csc", "csf", "dense", None}: - raise ValueError(f"{format} format not supported.") - _validate_device(device) - tensor = ( - obj - if isinstance(obj, Tensor) - else Tensor(obj, fill_value=fill_value, copy=copy) - ) - if format is not None: - if copy is False: - raise ValueError( - "Unable to avoid copy while creating an array as requested." - ) - order = tensor.get_order() - if format == "coo": - storage = Storage(SparseCOO(tensor.ndim, Element(tensor.fill_value)), order) - elif format == "csr": - storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (2, 1)) - elif format == "csc": - storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (1, 2)) - elif format == "csf": - storage = Element(tensor.fill_value) - for _ in range(tensor.ndim - 1): - storage = SparseList(storage) - storage = Storage(Dense(storage), order) - elif format == "dense": - storage = DenseStorage(tensor.ndim, tensor.dtype, order) - tensor = tensor.to_storage(storage) - - if dtype is not None: - return astype(tensor, dtype, copy=copy) - return tensor - - -def reshape( - x: Tensor, /, shape: tuple[int, ...], *, copy: bool | None = None -) -> Tensor: - if copy is False: - raise ValueError("Unable to avoid copy during reshape.") - # TODO: https://github.com/finch-tensor/Finch.jl/issues/743 - # Revert to `jl.reshape` implementation once aforementioned - # issue is solved. - warnings.warn( - "`reshape` densified the input tensor.", PerformanceWarning, stacklevel=2 - ) - arr = x.todense() - arr = arr.reshape(shape) - return Tensor(arr) - - -def _validate_device(device: Device) -> None: - if device not in {"cpu", None}: - raise ValueError( - f'Device not understood. Only "cpu" is allowed, but received: {device}' - ) +# def element_type(self) -> Any: +# return jl.eltype(self.jltype) + +# # TODO: implement later +# def shape_type(self) -> tuple[type, ...]: +# ... + +# def __call__(self, shape: tuple) -> Tensor: +# ... + +# def from_numpy(self, arr: np.ndarray) -> Tensor: +# ... + +# class FinchJLTensor(_Display, EagerTensor): +# """ +# A wrapper class for Finch.Tensor and Finch.SwizzleArray. + +# Constructors +# ------------ +# FinchJLTensor(scipy.sparse.spmatrix) +# Construct a Tensor out of a `scipy.sparse` object. Supported formats are: `COO`, +# `CSC`, and `CSR`. +# FinchJLTensor(numpy.ndarray) +# Construct a Tensor out of a NumPy array object. This is a no-copy operation. +# FinchJLTensor(Storage) +# Initialize a Tensor with a `storage` description. `storage` can already hold +# data. +# FinchJLTensor(julia_object) +# Tensor created from a compatible raw Julia object. Must be a `Tensor`. +# This is a no-copy operation. + +# Parameters +# ---------- +# obj : np.ndarray or scipy.sparse or Storage or Finch.Tensor +# Input to construct a Tensor. It's a no-copy operation of for NumPy and +# SciPy input. For Storage it's levels' description with order. The order +# numbers the dimensions from the fastest to slowest. The leaf nodes have +# mode `0` and the root node has mode `n-1`. If the tensor was square of +# size `N`, then `N .^ order == strides`. Available options are "C" +# (row-major), "F" (column-major), or a custom order. Default: row-major. +# fill_value : np.number, optional +# Only used when `numpy.ndarray` or `scipy.sparse` is passed. +# copy : bool, optional +# If ``True``, then the object is copied. If ``None`` then the object is +# copied only if needed. For ``False`` it raises a ``ValueError`` if a +# copy cannot be avoided. Default: ``None``. + +# Returns +# ------- +# FinchJLTensor +# Python wrapper for Finch.jl `Tensor`. + +# Examples +# -------- +# >>> import numpy as np +# >>> import finch +# >>> arr2d = np.arange(6).reshape((2, 3)) +# >>> t1 = finch.FinchJLTensor(arr2d) +# >>> t1.todense() +# array([[0, 1, 2], +# [3, 4, 5]]) +# >>> np.shares_memory(t1.todense(), arr2d) +# True +# >>> storage = finch.Storage( +# ... finch.Dense(finch.SparseList(finch.Element(1))), order="C" +# ... ) +# >>> t2 = t1.to_storage(storage) +# >>> t2.todense() +# array([[0, 1, 2], +# [3, 4, 5]]) +# """ + +# def __init__( +# self, +# obj: np.ndarray | spmatrix | Storage | JuliaObj, +# /, +# *, +# fill_value: np.number | None = None, +# copy: bool | None = None, +# ): +# if isinstance(obj, int | float | complex | bool | list): +# if copy is False: +# raise ValueError( +# "copy=False isn't supported for scalar inputs and Python lists" +# ) +# obj = np.asarray(obj) +# if fill_value is None: +# fill_value = 0.0 + +# if _is_scipy_sparse_obj(obj): # scipy constructor +# jl_data = self._from_scipy_sparse(obj, fill_value=fill_value, copy=copy) +# self._obj = jl_data +# elif isinstance(obj, np.ndarray): # numpy constructor +# jl_data = self._from_numpy(obj, fill_value=fill_value, copy=copy) +# self._obj = jl_data +# elif isinstance(obj, Storage): # from-storage constructor +# if copy: +# self._raise_julia_copy_not_supported() +# order = self.preprocess_order( +# obj.order, self.get_lvl_ndim(obj.levels_descr._obj) +# ) +# self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order) +# elif jl.isa(obj, jl.Finch.Tensor): # raw-Julia-object constructors +# if copy: +# self._raise_julia_copy_not_supported() +# self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1))) +# elif jl.isa(obj, jl.Finch.Tensor): +# if copy: +# self._raise_julia_copy_not_supported() +# self._obj = obj +# elif isinstance(obj, FinchJLTensor): +# self._obj = obj._obj +# else: +# raise ValueError( +# "Either scalar, numpy, scipy.sparse or a raw julia object should " +# f"be provided. Found: {type(obj)}" +# ) + +# @property +# def element_type(self): +# return jl.eltype(self._obj.body) + +# @property +# def dtype(self) -> DType: +# return jl.eltype(self._obj.body) + +# @property +# def ndim(self) -> int: +# return jl.ndims(self._obj) + +# @property +# def shape(self) -> tuple[int, ...]: +# return jl.size(self._obj) + +# @property +# def size(self) -> int: +# return np.prod(self.shape) + +# @property +# def fill_value(self) -> np.number: +# return jl.fill_value(self._obj) + +# @property +# def _is_dense(self) -> bool: +# lvl = self._obj.body.lvl +# for _ in self.shape: +# if not jl.isa(lvl, jl.Finch.Dense): +# return False +# lvl = lvl.lvl +# return True + +# @property +# def _order(self) -> tuple[int, ...]: +# return jl.typeof(self._obj).parameters[1] + +# @property +# def mT(self) -> Tensor: +# axes = list(range(self.ndim)) +# axes[-2], axes[-1] = axes[-1], axes[-2] +# axes = tuple(axes) +# return self.permute_dims(axes) + +# @property +# def device(self) -> str: +# return "cpu" + +# def to_device( +# self, device: Device, /, *, stream: int | Any | None = None +# ) -> Tensor: +# if device != "cpu": +# raise ValueError("Only `device='cpu'` is supported.") + +# return self + +# @classmethod +# def get_lvl_ndim(cls, lvl: JuliaObj) -> int: +# ndim = 0 +# while True: +# ndim += 1 +# lvl = lvl.lvl +# if jl.isa(lvl, jl.Finch.Element): +# break +# return ndim + +# def todense(self) -> np.ndarray: +# obj = self._obj + +# if self._is_dense: +# # don't materialize a dense finch tensor +# shape = jl.size(obj.body) +# dense_tensor = obj.body.lvl +# else: +# # create materialized dense array +# shape = jl.size(obj) +# dense_lvls = jl.Element(jc.convert(self.dtype, jl.fill_value(obj))) +# for _ in range(self.ndim): +# dense_lvls = jl.Dense(dense_lvls) +# dense_tensor = jl.Tensor(dense_lvls, obj).lvl # materialize + +# for _ in range(self.ndim): +# dense_tensor = dense_tensor.lvl + +# result = np.asarray(jl.reshape(dense_tensor.val, shape)) +# return result.transpose(self.get_order()) if self._is_dense else result + +# #TODO: Do we need? +# def permute_dims(self, axes: tuple[int, ...]) -> Tensor: +# axes = tuple(i + 1 for i in axes) +# new_obj = jl.permutedims(self._obj, axes) +# return Tensor(new_obj) + +# def to_storage(self, storage: Storage) -> Tensor: +# return Tensor(self._from_other_tensor(self, storage=storage)) + +# @classmethod +# def _from_other_tensor(cls, tensor: Tensor, storage: Storage) -> JuliaObj: +# order = cls.preprocess_order(storage.order, tensor.ndim) +# result = jl.copyto_b( +# jl.swizzle(jl.Tensor(storage.levels_descr._obj), *order), tensor._obj +# ) +# return jl.dropfills(result) if tensor._is_dense else result + +# @classmethod +# def _from_numpy( +# cls, arr: np.ndarray, fill_value: np.number, copy: bool | None = None +# ) -> JuliaObj: +# if copy: +# arr = arr.copy() +# order_char = "F" if np.isfortran(arr) else "C" +# order = cls.preprocess_order(order_char, arr.ndim) +# inv_order = tuple(i - 1 for i in jl.invperm(order)) + +# dtype = arr.dtype.type +# if ( +# dtype == np.bool_ +# ): # Fails with: Finch currently only supports isbits defaults +# dtype = jl_dtypes.bool +# fill_value = dtype(fill_value) +# lvl = Element(fill_value, arr.reshape(-1, order=order_char)) +# for i in inv_order: +# lvl = Dense(lvl, arr.shape[i]) +# return jl.swizzle(jl.Tensor(lvl._obj), *order) + +# @classmethod +# def from_scipy_sparse( +# cls, +# x, +# fill_value: np.number | None = None, +# copy: bool | None = None, +# ) -> Tensor: +# if not _is_scipy_sparse_obj(x): +# raise ValueError("{x} is not a SciPy sparse object.") +# return Tensor(x, fill_value=fill_value, copy=copy) + +# @classmethod +# def _from_scipy_sparse( +# cls, +# x, +# *, +# fill_value: np.number | None = None, +# copy: bool | None = None, +# ) -> JuliaObj: +# if copy is False and not ( +# x.format in ("coo", "csr", "csc") and x.has_canonical_format +# ): +# raise ValueError( +# "Unable to avoid copy while creating an array as requested." +# ) +# if x.format not in ("coo", "csr", "csc"): +# x = x.asformat("coo") +# if copy: +# x = x.copy() +# if not x.has_canonical_format: +# x.sum_duplicates() +# assert x.has_canonical_format + +# if x.format == "coo": +# return cls.construct_coo_jl_object( +# coords=(x.col, x.row), +# data=x.data, +# shape=x.shape[::-1], +# order=Tensor.row_major, +# fill_value=fill_value, +# ) +# if x.format == "csc": +# return cls.construct_csc_jl_object( +# arg=(x.data, x.indices, x.indptr), +# shape=x.shape, +# fill_value=fill_value, +# ) +# if x.format == "csr": +# return cls.construct_csr_jl_object( +# arg=(x.data, x.indices, x.indptr), +# shape=x.shape, +# fill_value=fill_value, +# ) +# raise ValueError(f"Unsupported SciPy format: {type(x)}") + +# @classmethod +# def construct_coo_jl_object( +# cls, coords, data, shape, order, fill_value=0.0 +# ) -> JuliaObj: +# assert len(coords) == 2 +# ndim = len(shape) +# order = cls.preprocess_order(order, ndim) + +# lvl = jl.Element(data.dtype.type(fill_value), data) +# ptr = jl.Vector[jl.Int]([1, len(data) + 1]) +# tbl = tuple(jl.PlusOneVector(arr) for arr in coords) + +# return jl.swizzle(jl.Tensor(jl.SparseCOO[ndim](lvl, shape, ptr, tbl)), *order) + +# @classmethod +# def construct_coo( +# cls, coords, data, shape, order=row_major, fill_value=0.0 +# ) -> Tensor: +# return Tensor( +# cls.construct_coo_jl_object(coords, data, shape, order, fill_value) +# ) + +# @staticmethod +# def _construct_compressed2d_jl_object( +# arg: TupleOf3Arrays, +# shape: tuple[int, ...], +# order: tuple[int, ...], +# fill_value: np.number = 0.0, +# ) -> JuliaObj: +# assert isinstance(arg, tuple) and len(arg) == 3 +# assert len(shape) == 2 + +# data, indices, indptr = arg +# dtype = data.dtype.type +# indices = jl.PlusOneVector(indices) +# indptr = jl.PlusOneVector(indptr) + +# lvl = jl.Element(dtype(fill_value), data) +# return jl.swizzle( +# jl.Tensor( +# jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1]) +# ), +# *order, +# ) + +# @classmethod +# def construct_csc_jl_object( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> JuliaObj: +# return cls._construct_compressed2d_jl_object( +# arg=arg, shape=shape, order=(1, 2), fill_value=fill_value +# ) + +# @classmethod +# def construct_csc( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> Tensor: +# return Tensor(cls.construct_csc_jl_object(arg, shape, fill_value)) + +# @classmethod +# def construct_csr_jl_object( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> JuliaObj: +# return cls._construct_compressed2d_jl_object( +# arg=arg, shape=shape[::-1], order=(2, 1), fill_value=fill_value +# ) + +# @classmethod +# def construct_csr( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> Tensor: +# return Tensor(cls.construct_csr_jl_object(arg, shape, fill_value)) + +# @staticmethod +# def construct_csf_jl_object( +# arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> JuliaObj: +# assert isinstance(arg, tuple) and len(arg) == 3 + +# data, indices_list, indptr_list = arg +# dtype = data.dtype.type + +# assert len(indices_list) == len(shape) - 1 +# assert len(indptr_list) == len(shape) - 1 + +# indices_list = [jl.PlusOneVector(i) for i in indices_list] +# indptr_list = [jl.PlusOneVector(i) for i in indptr_list] + +# lvl = jl.Element(dtype(fill_value), data) +# for size, indices, indptr in zip( +# shape[:-1], indices_list, indptr_list, strict=False +# ): +# lvl = jl.SparseList(lvl, size, indptr, indices) + +# return jl.swizzle( +# jl.Tensor(jl.Dense(lvl, shape[-1])), *range(1, len(shape) + 1) +# ) + +# @classmethod +# def construct_csf( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> Tensor: +# return Tensor(cls.construct_csf_jl_object(arg, shape, fill_value)) + +# def to_scipy_sparse(self, accept_fv=None): +# import scipy.sparse as sp + +# if accept_fv is None: +# accept_fv = [0] +# elif not isinstance(accept_fv, Iterable): +# accept_fv = [accept_fv] + +# if self.ndim != 2: +# raise ValueError( +# "Can only convert a 2-dimensional array to a Scipy sparse matrix." +# ) +# if not builtins.any(_eq_scalars(self.fill_value, fv) for fv in accept_fv): +# raise ValueError( +# f"Can only convert arrays with {accept_fv} fill-values " +# "to a Scipy sparse matrix." +# ) +# order = self.get_order() +# body = self._obj.body + +# if str(jl.typeof(body.lvl).name.name) == "SparseCOOLevel": +# data = np.asarray(body.lvl.lvl.val) +# coords = body.lvl.tbl +# row, col = coords[::-1] if order == (1, 0) else coords +# row, col = np.asarray(row) - 1, np.asarray(col) - 1 +# return sp.coo_matrix((data, (row, col)), shape=self.shape) + +# if ( +# str(jl.typeof(body.lvl).name.name) == "DenseLevel" +# and str(jl.typeof(body.lvl.lvl).name.name) == "SparseListLevel" +# ): +# data = np.asarray(body.lvl.lvl.lvl.val) +# indices = np.asarray(body.lvl.lvl.idx) - 1 +# indptr = np.asarray(body.lvl.lvl.ptr) - 1 +# sp_class = sp.csr_matrix if order == (1, 0) else sp.csc_matrix +# return sp_class((data, indices, indptr), shape=self.shape) +# if ( +# jl.typeof(body.lvl).name.name in sparse_formats_names +# or jl.typeof(body.lvl.lvl).name.name in sparse_formats_names +# ): +# storage = Storage(SparseCOO(self.ndim, Element(self.fill_value)), order) +# return self.to_storage(storage).to_scipy_sparse() +# raise ValueError("Tensor can't be converted to scipy.sparse object.") + +# @staticmethod +# def _raise_julia_copy_not_supported() -> None: +# raise ValueError("copy=True isn't supported for Julia object inputs") + +# def __array_namespace__(self, *, api_version: str | None = None) -> Any: +# if api_version is None: +# api_version = "2024.12" + +# if api_version not in {"2021.12", "2022.12", "2023.12", "2024.12"}: +# raise ValueError(f'"{api_version}" Array API version not supported.') +# import finch + +# return finch + + +# def random(shape, density=0.01, random_state=None): +# args = [*shape, density] +# if random_state is not None: +# if isinstance(random_state, np.random.Generator): +# seed = random_state.integers(np.iinfo(np.int32).max) +# else: +# seed = random_state +# rng = jl.Random.default_rng() +# jl.Random.seed_b(rng, seed) +# args = [rng] + args +# return Tensor(jl.fsprand(*args)) + + +# def asarray( +# obj, +# /, +# *, +# dtype: DType | None = None, +# format: str | None = None, +# fill_value: np.number | None = None, +# device: Device | None = None, +# copy: bool | None = None, +# ) -> Tensor: +# if format not in {"coo", "csr", "csc", "csf", "dense", None}: +# raise ValueError(f"{format} format not supported.") +# _validate_device(device) +# tensor = ( +# obj +# if isinstance(obj, Tensor) +# else Tensor(obj, fill_value=fill_value, copy=copy) +# ) +# if format is not None: +# if copy is False: +# raise ValueError( +# "Unable to avoid copy while creating an array as requested." +# ) +# order = tensor.get_order() +# if format == "coo": +# storage = Storage(SparseCOO(tensor.ndim, Element(tensor.fill_value)), order) +# elif format == "csr": +# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (2, 1)) +# elif format == "csc": +# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (1, 2)) +# elif format == "csf": +# storage = Element(tensor.fill_value) +# for _ in range(tensor.ndim - 1): +# storage = SparseList(storage) +# storage = Storage(Dense(storage), order) +# elif format == "dense": +# storage = DenseStorage(tensor.ndim, tensor.dtype, order) +# tensor = tensor.to_storage(storage) + +# if dtype is not None: +# return astype(tensor, dtype, copy=copy) +# return tensor + + +# def reshape( +# x: Tensor, /, shape: tuple[int, ...], *, copy: bool | None = None +# ) -> Tensor: +# if copy is False: +# raise ValueError("Unable to avoid copy during reshape.") +# # TODO: https://github.com/finch-tensor/Finch.jl/issues/743 +# # Revert to `jl.reshape` implementation once aforementioned +# # issue is solved. +# warnings.warn( +# "`reshape` densified the input tensor.", PerformanceWarning, stacklevel=2 +# ) +# arr = x.todense() +# arr = arr.reshape(shape) +# return Tensor(arr) + + +# def _validate_device(device: Device) -> None: +# if device not in {"cpu", None}: +# raise ValueError( +# f'Device not understood. Only "cpu" is allowed, but received: {device}' +# ) diff --git a/src/finch/typing.py b/src/finch/typing.py index 9c96fe5..fa305c9 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -1,17 +1,15 @@ -from typing import Any, Literal +from typing import Literal, Any import numpy as np import juliacall as jc -OrderType = Literal["C", "F"] | tuple[int, ...] | None - TupleOf3Arrays = tuple[np.ndarray, np.ndarray, np.ndarray] +spmatrix = Any + JuliaObj = jc.AnyValue DType = jc.AnyValue # represents jl.DataType -spmatrix = Any - Device = Literal["cpu"] | None diff --git a/tests/test_compiler.py b/tests/test_compiler.py new file mode 100644 index 0000000..68783f9 --- /dev/null +++ b/tests/test_compiler.py @@ -0,0 +1,192 @@ +import pytest +import numpy as np +from finchlite.finch_notation.nodes import ( + Module, + Function, + Variable, + Block, + Assign, + Call, + Literal, + Update, + Unpack, + Slot, + Loop, + Increment, + Access, + Unwrap, + Freeze, + Read, + Repack, + Return, + Declare, +) + +import operator +from finchlite import ftype +from finchlite.algebra import overwrite, promote_min +from finchlite.compile import ExtentFType, dimension, BufferizedNDArray +from finchlite.codegen import NumpyBuffer + +from finch.compiler import FinchJLCompiler + +# Dummy data to obtain the bufferized ND array type +a = np.zeros(dtype=np.float64, shape=(3, 3)) +a_format = ftype(BufferizedNDArray.from_numpy(a)) + +@pytest.mark.parametrize( + "finch_ntn, julia_code", + [ + ( + Module( + ( + Function( + Variable("matmul", a_format), + ( + Variable("C", a_format), + Variable("A", a_format), + Variable("B", a_format), + ), + Block( + ( + Assign( + Variable("m", ExtentFType(np.int64, np.int64)), + Call( + Literal(dimension), + (Variable("A", a_format), Literal(0)), + ), + ), + Assign( + Variable("n", ExtentFType(np.int64, np.int64)), + Call( + Literal(dimension), + (Variable("B", a_format), Literal(1)), + ), + ), + Assign( + Variable("p", ExtentFType(np.int64, np.int64)), + Call( + Literal(dimension), + (Variable("A", a_format), Literal(1)), + ), + ), + Unpack(Slot("A_", a_format), Variable("A", a_format)), + Unpack(Slot("B_", a_format), Variable("B", a_format)), + Unpack(Slot("C_", a_format), Variable("C", a_format)), + Declare( + Slot("C_", a_format), + Literal(0.0), + Literal(operator.add), + ( + Variable("m", ExtentFType(np.int64, np.int64)), + Variable("n", ExtentFType(np.int64, np.int64)), + ), + ), + Loop( + Variable("i", np.int64), + Variable("m", ExtentFType(np.int64, np.int64)), + Loop( + Variable("k", np.int64), + Variable("p", ExtentFType(np.int64, np.int64)), + Loop( + Variable("j", np.int64), + Variable( + "n", ExtentFType(np.int64, np.int64) + ), + Block( + ( + Assign( + Variable("a_ik", np.float64), + Unwrap( + Access( + Slot("A_", a_format), + Read(), + ( + Variable( + "i", np.int64 + ), + Variable( + "k", np.int64 + ), + ), + ) + ), + ), + Assign( + Variable("b_kj", np.float64), + Unwrap( + Access( + Slot("B_", a_format), + Read(), + ( + Variable( + "k", np.int64 + ), + Variable( + "j", np.int64 + ), + ), + ) + ), + ), + Assign( + Variable("c_ij", np.float64), + Call( + Literal(operator.mul), + ( + Variable( + "a_ik", np.float64 + ), + Variable( + "b_kj", np.float64 + ), + ), + ), + ), + Increment( + Access( + Slot("C_", a_format), + Update( + Literal(operator.add) + ), + ( + Variable("i", np.int64), + Variable("j", np.int64), + ), + ), + Variable("c_ij", np.float64), + ), + ) + ), + ), + ), + ), + Freeze(Slot("C_", a_format), Literal(operator.add)), + Repack(Slot("C_", a_format), Variable("C", a_format)), + Return(Variable("C", a_format)), + ) + ), + ), + ) + ), + """function matmul(C,A,B) + C .= 0 + for i = _ + for k = _ + for j = _ + a_ik = A[i,k] + b_kj = B[k,j] + c_ij = a_ik * b_kj + C[i,j] = c_ij + end + end + end + return C +end""", + ) + ], +) +def test_finchjl_compiler(finch_ntn: Module, julia_code): + compiler = FinchJLCompiler() + library = compiler(finch_ntn) + assert getattr(library, finch_ntn.children[0].name.name).jl_code == julia_code From 22b2f71bc6bf782b9b1c836066b6ee772fd2fc41 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 10 Feb 2026 23:41:55 -0500 Subject: [PATCH 04/81] wip: added match cases for most finch-tensor-lite finch-ntn nodes --- src/finch/compiler.py | 155 +++++++++++++++++++++++++++++------------- 1 file changed, 106 insertions(+), 49 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 10281cc..7ea8f08 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,10 +1,15 @@ from finchlite.compile import NotationCompiler from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary import finchlite.finch_notation.nodes as ntn +from finchlite.compile import dimension from typing import Any +import operator + from juliacall import Main as jl +ops_map = {operator.add: "+", operator.mul: "*"} + class FinchJLKernel(AssemblyKernel): def __init__(self, func_name, jl_code): @@ -31,61 +36,113 @@ def __getattr__(self, name: str) -> FinchJLKernel: class FinchJLGenerator: + def __init__(self): + self.pack_dict = {} + def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: + self.pack_dict.clear() + return self.generate_julia(prgm) + + def generate_julia(self, prgm, nestingLvl=0): match prgm: - case ntn.Literal(val): - ... - # What this? - case ntn.Value(ex, type_): - ... - case ntn.Variable(name, type_): - ... + case ntn.Function(name, args, body): + body_str = self.generate_julia(body, nestingLvl + 1) + return f"function {name}\n{body_str}\nend" + + case ntn.Block(bodies): + body_str = "" + tab_str = {" " * nestingLvl} + for body in bodies: + body_str += f"{tab_str}{self.generate_julia(body, nestingLvl)}\n" + + case ntn.Assign(lhs, rhs): + # TODO: Can we make this better? + # Special condition to ignore all assigns associated with + # finding loop bounds + if isinstance(rhs, ntn.Call) and rhs.op.val == dimension: + return "" + return f"{self.generate_julia(lhs, nestingLvl)} = {self.generate_julia(body, nestingLvl)}" + + case ntn.Declare(tns, init, op, shape): + # TODO: what is the purpose of op here + return f"@finch {self.generate_julia(tns, nestingLvl)} .= {self.generate_julia(init, nestingLvl)}" + + case ntn.Return(val): + return f"return {self.generate_julia(val, nestingLvl)}" + + case ntn.Loop(idx, _, body): + tab_str = " " * nestingLvl + loop_body = self.generate_julia(body, nestingLvl + 1) + return f"for {idx.name} = _\n{loop_body}\n{tab_str}end" + + case ntn.Access(tns, _, idxs): + tns_str = self.generate_julia(tns, nestingLvl) + idx_str = [self.generate_julia(idx, nestingLvl) for idx in idxs].join( + "," + ) + return f"{tns_str}[{idx_str}]" + case ntn.Call(op, args): - ... - case ntn.Dimension(tns, r): - ... - case ntn.Access(tns, mode, idxs): - ... - case ntn.Read(): - ... - case ntn.AccessMode(): - ... - case ntn.Update(op): - ... - case ntn.Increment(lhs,rhs): - ... - case ntn.Unwrap(arg): - ... - case ntn.Cached(arg,ref): - ... - case ntn.Loop(idx, ext, body): - ... + arg_str = [self.generate_julia(arg, nestingLvl) for arg in args].join( + "," + ) + return f"{ops_map[op.val]}({arg_str})" + case ntn.If(cond, body): - ... + tab_str = " " * nestingLvl + cond_str = self.generate_julia(cond, nestingLvl) + body_str = self.generate_julia(body, nestingLvl + 1) + return f"if {cond_str}\n{body_str}\n{tab_str}end" + case ntn.IfElse(cond, then_body, else_body): - ... - case ntn.Assign(lhs,rhs): - ... - case ntn.Stack(obj,type): - ... - case ntn.Slot(name, type): - ... + cond_str = self.generate_julia(cond, nestingLvl) + then_body_str = self.generate_julia(then_body, nestingLvl + 1) + else_body_str = self.generate_julia(else_body, nestingLvl + 1) + return f"if {cond_str}\n{then_body_str}\n{tab_str}else\n{else_body_str}\n{tab_str}end" + + case ntn.Increment(lhs, rhs): + lhs_str = self.generate_julia(lhs, nestingLvl) + rhs_str = self.generate_julia(rhs, nestingLvl) + + # TODO: Is this the correct assumption to make + if not ( + isinstance(lhs, ntn.Access) and isinstance(lhs.mode, ntn.Update) + ): + raise Exception("Increment expects the lhs to be an access") + + return f"{lhs_str} {ops_map[lhs.mode.op.val]}= {rhs_str}" + + case ntn.Unwrap(arg): + return self.generate_julia(arg, nestingLvl) + case ntn.Unpack(lhs, rhs): - ... - case ntn.Repack(val, obj): - ... - case ntn.Declare(tns, init, op, shape): - ... - case ntn.Freeze(tns, op): - ... - case ntn.Thaw(tns, op): - ... - case ntn.Block(bodies): - ... - case ntn.Function(name,args,body): - ... - case ntn.Return(val): - ... + # TODO: Is this the right assumption to make + if not isinstance(rhs, ntn.Variable): + raise Exception("The unpack was not called with variable as RHS.") + self.pack_dict[lhs.name] = rhs.name + return "" + + case ntn.Repack(val, _): + self.pack_dict.pop(val.name) + return "" + + case ntn.Freeze(_, _): + return "" + + case ntn.Slot(name): + if name not in self.pack_dict: + raise Exception(f"{name} Slot does not exist in registry.") + return self.pack_dict[name] + + case ntn.Literal(val): + return str(val) + + case ntn.Variable(name, _): + return name + + # TODO: Cached, Dimension, Thaw, Stack, Value are unimplemented. + case _: + raise Exception(f"Unhandled node type: {type(prgm)}") class FinchJLCompiler(NotationCompiler): From 9ee64427c711bdc6f85b50264bd1b399aa41bbeb Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 11 Feb 2026 00:43:40 -0500 Subject: [PATCH 05/81] feat: added support for compiling to julia string --- src/finch/compiler.py | 63 +++++++++++++++------ tests/test_compiler.py | 124 ++++++++++++++++++++--------------------- 2 files changed, 105 insertions(+), 82 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 7ea8f08..a47835a 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -15,7 +15,9 @@ class FinchJLKernel(AssemblyKernel): def __init__(self, func_name, jl_code): self.jl_code = jl_code self.func_name = func_name - jl.seval(jl_code) + + print(jl_code) + # jl.seval(jl_code) # TODO: Switch back to (self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...] def __call__(self, *args: tuple[Any, ...]): @@ -24,7 +26,7 @@ def __call__(self, *args: tuple[Any, ...]): argList.append(f"arg{len(argList)}") setattr(jl, argList[-1], arg) - jl.seval(f"{self.func_name}({argList.join(',')})") + jl.seval(f"{self.func_name}({','.join(argList)})") class FinchJLLibrary(AssemblyLibrary): @@ -32,28 +34,35 @@ def __init__(self, kernel_dict): self.kernel_dict = kernel_dict def __getattr__(self, name: str) -> FinchJLKernel: - return self.kernels[name] + return self.kernel_dict[name] class FinchJLGenerator: def __init__(self): self.pack_dict = {} + self.in_finch_block = False def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: self.pack_dict.clear() + self.in_finch_block = False return self.generate_julia(prgm) def generate_julia(self, prgm, nestingLvl=0): match prgm: case ntn.Function(name, args, body): body_str = self.generate_julia(body, nestingLvl + 1) - return f"function {name}\n{body_str}\nend" + arg_str = ",".join( + [self.generate_julia(arg, nestingLvl) for arg in args] + ) + return f"function {name}({arg_str})\n{body_str}end" case ntn.Block(bodies): body_str = "" - tab_str = {" " * nestingLvl} for body in bodies: - body_str += f"{tab_str}{self.generate_julia(body, nestingLvl)}\n" + curr_body_str = self.generate_julia(body, nestingLvl) + if curr_body_str != "": + body_str += f"{curr_body_str}\n" + return body_str case ntn.Assign(lhs, rhs): # TODO: Can we make this better? @@ -61,30 +70,47 @@ def generate_julia(self, prgm, nestingLvl=0): # finding loop bounds if isinstance(rhs, ntn.Call) and rhs.op.val == dimension: return "" - return f"{self.generate_julia(lhs, nestingLvl)} = {self.generate_julia(body, nestingLvl)}" + + tab_str = " " * nestingLvl + return f"{tab_str}{self.generate_julia(lhs, nestingLvl)} = {self.generate_julia(rhs, nestingLvl)}" case ntn.Declare(tns, init, op, shape): # TODO: what is the purpose of op here - return f"@finch {self.generate_julia(tns, nestingLvl)} .= {self.generate_julia(init, nestingLvl)}" + tab_str = " " * nestingLvl + return f"{tab_str}@finch {self.generate_julia(tns, nestingLvl)} .= {self.generate_julia(init, nestingLvl)}" case ntn.Return(val): - return f"return {self.generate_julia(val, nestingLvl)}" + tab_str = " " * nestingLvl + return f"{tab_str}return {self.generate_julia(val, nestingLvl)}" case ntn.Loop(idx, _, body): tab_str = " " * nestingLvl - loop_body = self.generate_julia(body, nestingLvl + 1) - return f"for {idx.name} = _\n{loop_body}\n{tab_str}end" + tab_str_1 = " " * (nestingLvl+1) + + is_outermost_loop = False + if self.in_finch_block is False: + is_outermost_loop = True + self.in_finch_block = True + loop_body = self.generate_julia(body, nestingLvl + 2) + else: + loop_body = self.generate_julia(body, nestingLvl + 1) + + if not is_outermost_loop: + return f"{tab_str}for {idx.name} = _\n{loop_body}{tab_str}end\n" + else: + self.in_finch_block = False + return f"{tab_str}@finch begin\n{tab_str_1}for {idx.name} = _\n{loop_body}{tab_str_1}end\n{tab_str}end" case ntn.Access(tns, _, idxs): tns_str = self.generate_julia(tns, nestingLvl) - idx_str = [self.generate_julia(idx, nestingLvl) for idx in idxs].join( - "," + idx_str = ",".join( + [self.generate_julia(idx, nestingLvl) for idx in idxs] ) return f"{tns_str}[{idx_str}]" case ntn.Call(op, args): - arg_str = [self.generate_julia(arg, nestingLvl) for arg in args].join( - "," + arg_str = ",".join( + [self.generate_julia(arg, nestingLvl) for arg in args] ) return f"{ops_map[op.val]}({arg_str})" @@ -92,15 +118,16 @@ def generate_julia(self, prgm, nestingLvl=0): tab_str = " " * nestingLvl cond_str = self.generate_julia(cond, nestingLvl) body_str = self.generate_julia(body, nestingLvl + 1) - return f"if {cond_str}\n{body_str}\n{tab_str}end" + return f"{tab_str}if {cond_str}\n{body_str}\n{tab_str}end" case ntn.IfElse(cond, then_body, else_body): cond_str = self.generate_julia(cond, nestingLvl) then_body_str = self.generate_julia(then_body, nestingLvl + 1) else_body_str = self.generate_julia(else_body, nestingLvl + 1) - return f"if {cond_str}\n{then_body_str}\n{tab_str}else\n{else_body_str}\n{tab_str}end" + return f"{tab_str}if {cond_str}\n{then_body_str}\n{tab_str}else\n{else_body_str}\n{tab_str}end" case ntn.Increment(lhs, rhs): + tab_str = " " * nestingLvl lhs_str = self.generate_julia(lhs, nestingLvl) rhs_str = self.generate_julia(rhs, nestingLvl) @@ -110,7 +137,7 @@ def generate_julia(self, prgm, nestingLvl=0): ): raise Exception("Increment expects the lhs to be an access") - return f"{lhs_str} {ops_map[lhs.mode.op.val]}= {rhs_str}" + return f"{tab_str}{lhs_str} {ops_map[lhs.mode.op.val]}= {rhs_str}" case ntn.Unwrap(arg): return self.generate_julia(arg, nestingLvl) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 68783f9..6c51fa5 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -34,6 +34,7 @@ a = np.zeros(dtype=np.float64, shape=(3, 3)) a_format = ftype(BufferizedNDArray.from_numpy(a)) + @pytest.mark.parametrize( "finch_ntn, julia_code", [ @@ -95,54 +96,6 @@ ), Block( ( - Assign( - Variable("a_ik", np.float64), - Unwrap( - Access( - Slot("A_", a_format), - Read(), - ( - Variable( - "i", np.int64 - ), - Variable( - "k", np.int64 - ), - ), - ) - ), - ), - Assign( - Variable("b_kj", np.float64), - Unwrap( - Access( - Slot("B_", a_format), - Read(), - ( - Variable( - "k", np.int64 - ), - Variable( - "j", np.int64 - ), - ), - ) - ), - ), - Assign( - Variable("c_ij", np.float64), - Call( - Literal(operator.mul), - ( - Variable( - "a_ik", np.float64 - ), - Variable( - "b_kj", np.float64 - ), - ), - ), - ), Increment( Access( Slot("C_", a_format), @@ -154,30 +107,73 @@ Variable("j", np.int64), ), ), - Variable("c_ij", np.float64), + Call( + Literal(operator.mul), + ( + Unwrap( + Access( + Slot( + "A_", + a_format, + ), + Read(), + ( + Variable( + "i", + np.int64, + ), + Variable( + "k", + np.int64, + ), + ), + ) + ), + Unwrap( + Access( + Slot( + "B_", + a_format, + ), + Read(), + ( + Variable( + "k", + np.int64, + ), + Variable( + "j", + np.int64, + ), + ), + ) + ), + ), + ), ), - ) + ), ), ), ), ), - Freeze(Slot("C_", a_format), Literal(operator.add)), - Repack(Slot("C_", a_format), Variable("C", a_format)), - Return(Variable("C", a_format)), - ) + Freeze(Slot("C_", a_format), Literal(operator.add)), + Repack( + Slot("C_", a_format), Variable("C", a_format) + ), + Return(Variable("C", a_format)), + ), + ), ), - ), - ) + ) ), """function matmul(C,A,B) - C .= 0 - for i = _ - for k = _ - for j = _ - a_ik = A[i,k] - b_kj = B[k,j] - c_ij = a_ik * b_kj - C[i,j] = c_ij + @finch C .= 0.0 + @finch begin + for i = _ + for k = _ + for j = _ + C[i,j] += *(A[i,k],B[k,j]) + end end end end From 35ab8b1f606a53e10e6d02db811abe7a382f895e Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 11 Feb 2026 11:59:53 -0500 Subject: [PATCH 06/81] wip: setting up FinchJLTensor --- src/finch/levels.py | 140 ++++---- src/finch/tensor-old.py | 577 +++++++++++++++++++++++++++++++++ src/finch/tensor.py | 688 +++++++--------------------------------- 3 files changed, 761 insertions(+), 644 deletions(-) create mode 100644 src/finch/tensor-old.py diff --git a/src/finch/levels.py b/src/finch/levels.py index f65a5ad..eb10eaa 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -1,110 +1,110 @@ -# from .julia import jl -# from .typing import DType, JuliaObj, OrderType +from .julia import jl +from .typing import DType, JuliaObj -# class _Display: -# _obj: JuliaObj +class _Display: + _obj: JuliaObj -# def __repr__(self): -# return jl.sprint(jl.show, self._obj) + def __repr__(self): + return jl.sprint(jl.show, self._obj) -# def __str__(self): -# return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) + def __str__(self): + return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) -# # LEVEL +# LEVEL -# class AbstractLevel(_Display): -# pass +class AbstractLevel(_Display): + pass -# # core levels +# core levels -# class Dense(AbstractLevel): -# def __init__(self, lvl, shape=None): -# args = [lvl._obj] -# if shape is not None: -# args.append(shape) -# self._obj = jl.Dense(*args) +class Dense(AbstractLevel): + def __init__(self, lvl, shape=None): + args = [lvl._obj] + if shape is not None: + args.append(shape) + self._obj = jl.Dense(*args) -# class Element(AbstractLevel): -# def __init__(self, fill_value, data=None): -# args = [fill_value] -# if data is not None: -# args.append(data) -# self._obj = jl.Element(*args) +class Element(AbstractLevel): + def __init__(self, fill_value, data=None): + args = [fill_value] + if data is not None: + args.append(data) + self._obj = jl.Element(*args) -# class Pattern(AbstractLevel): -# def __init__(self): -# self._obj = jl.Pattern() +class Pattern(AbstractLevel): + def __init__(self): + self._obj = jl.Pattern() -# # advanced levels +# advanced levels -# class SparseList(AbstractLevel): -# def __init__(self, lvl): -# self._obj = jl.SparseList(lvl._obj) +class SparseList(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.SparseList(lvl._obj) -# class SparseByteMap(AbstractLevel): -# def __init__(self, lvl): -# self._obj = jl.SparseByteMap(lvl._obj) +class SparseByteMap(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.SparseByteMap(lvl._obj) -# class RepeatRLE(AbstractLevel): -# def __init__(self, lvl): -# self._obj = jl.RepeatRLE(lvl._obj) +class RepeatRLE(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.RepeatRLE(lvl._obj) -# class SparseVBL(AbstractLevel): -# def __init__(self, lvl): -# self._obj = jl.SparseVBL(lvl._obj) +class SparseVBL(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.SparseVBL(lvl._obj) -# class SparseCOO(AbstractLevel): -# def __init__(self, ndim, lvl): -# self._obj = jl.SparseCOO[ndim](lvl._obj) +class SparseCOO(AbstractLevel): + def __init__(self, ndim, lvl): + self._obj = jl.SparseCOO[ndim](lvl._obj) -# class SparseHash(AbstractLevel): -# def __init__(self, ndim, lvl): -# self._obj = jl.SparseHash[ndim](lvl._obj) +class SparseHash(AbstractLevel): + def __init__(self, ndim, lvl): + self._obj = jl.SparseHash[ndim](lvl._obj) -# sparse_formats_names = ( -# "SparseList", -# "Sparse", -# "SparseHash", -# "SparseCOO", -# "SparseRLE", -# "SparseVBL", -# "SparseBand", -# "SparsePoint", -# "SparseInterval", -# ) +sparse_formats_names = ( + "SparseList", + "Sparse", + "SparseHash", + "SparseCOO", + "SparseRLE", + "SparseVBL", + "SparseBand", + "SparsePoint", + "SparseInterval", +) -# # STORAGE +# STORAGE -# class Storage: -# def __init__(self, levels_descr: AbstractLevel, order: OrderType = None): -# self.levels_descr = levels_descr -# self.order = order if order is not None else "C" +class Storage: + def __init__(self, levels_descr: AbstractLevel, order: OrderType = None): + self.levels_descr = levels_descr + self.order = order if order is not None else "C" -# def __str__(self) -> str: -# return f"Storage(lvl={str(self.levels_descr)}, order={self.order})" + def __str__(self) -> str: + return f"Storage(lvl={str(self.levels_descr)}, order={self.order})" -# class DenseStorage(Storage): -# def __init__(self, ndim: int, dtype: DType, order: OrderType = None): -# lvl = Element(dtype(0)) -# for _ in range(ndim): -# lvl = Dense(lvl) +class DenseStorage(Storage): + def __init__(self, ndim: int, dtype: DType, order: OrderType = None): + lvl = Element(dtype(0)) + for _ in range(ndim): + lvl = Dense(lvl) -# super().__init__(levels_descr=lvl, order=order) + super().__init__(levels_descr=lvl, order=order) diff --git a/src/finch/tensor-old.py b/src/finch/tensor-old.py new file mode 100644 index 0000000..afdffb1 --- /dev/null +++ b/src/finch/tensor-old.py @@ -0,0 +1,577 @@ +# from __future__ import annotations + +# import builtins +# import warnings +# from collections.abc import Callable, Iterable +# from typing import Any, Literal + +# import numpy as np +# from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple + +# from . import dtypes as jl_dtypes +# from .errors import PerformanceWarning +# from .julia import jc, jl +# from .levels import ( +# Dense, +# DenseStorage, +# Element, +# SparseCOO, +# SparseList, +# Storage, +# _Display, +# sparse_formats_names, +# ) +# from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix +# from finchlite import Tensor, TensorFType, EagerTensor + + +# class FinchJLTensorFType(TensorFType): +# def __init__(self, jltype): +# # Julia type associated with the tensor +# self.jltype = jltype + +# def ndims(self) -> np.intp: +# return np.intp(jl.ndims(self.jltype)) + +# def fill_value(self) -> Any: +# return jl.fill_value(self.jltype) + +# def element_type(self) -> Any: +# return jl.eltype(self.jltype) + +# # TODO: implement later +# def shape_type(self) -> tuple[type, ...]: +# ... + +# def __call__(self, shape: tuple) -> Tensor: +# ... + +# def from_numpy(self, arr: np.ndarray) -> Tensor: +# ... + +# class FinchJLTensor(_Display, EagerTensor): +# """ +# A wrapper class for Finch.Tensor and Finch.SwizzleArray. + +# Constructors +# ------------ +# FinchJLTensor(scipy.sparse.spmatrix) +# Construct a Tensor out of a `scipy.sparse` object. Supported formats are: `COO`, +# `CSC`, and `CSR`. +# FinchJLTensor(numpy.ndarray) +# Construct a Tensor out of a NumPy array object. This is a no-copy operation. +# FinchJLTensor(Storage) +# Initialize a Tensor with a `storage` description. `storage` can already hold +# data. +# FinchJLTensor(julia_object) +# Tensor created from a compatible raw Julia object. Must be a `Tensor`. +# This is a no-copy operation. + +# Parameters +# ---------- +# obj : np.ndarray or scipy.sparse or Storage or Finch.Tensor +# Input to construct a Tensor. It's a no-copy operation of for NumPy and +# SciPy input. For Storage it's levels' description with order. The order +# numbers the dimensions from the fastest to slowest. The leaf nodes have +# mode `0` and the root node has mode `n-1`. If the tensor was square of +# size `N`, then `N .^ order == strides`. Available options are "C" +# (row-major), "F" (column-major), or a custom order. Default: row-major. +# fill_value : np.number, optional +# Only used when `numpy.ndarray` or `scipy.sparse` is passed. +# copy : bool, optional +# If ``True``, then the object is copied. If ``None`` then the object is +# copied only if needed. For ``False`` it raises a ``ValueError`` if a +# copy cannot be avoided. Default: ``None``. + +# Returns +# ------- +# FinchJLTensor +# Python wrapper for Finch.jl `Tensor`. + +# Examples +# -------- +# >>> import numpy as np +# >>> import finch +# >>> arr2d = np.arange(6).reshape((2, 3)) +# >>> t1 = finch.FinchJLTensor(arr2d) +# >>> t1.todense() +# array([[0, 1, 2], +# [3, 4, 5]]) +# >>> np.shares_memory(t1.todense(), arr2d) +# True +# >>> storage = finch.Storage( +# ... finch.Dense(finch.SparseList(finch.Element(1))), order="C" +# ... ) +# >>> t2 = t1.to_storage(storage) +# >>> t2.todense() +# array([[0, 1, 2], +# [3, 4, 5]]) +# """ + +# def __init__( +# self, +# obj: np.ndarray | spmatrix | Storage | JuliaObj, +# /, +# *, +# fill_value: np.number | None = None, +# copy: bool | None = None, +# ): +# if isinstance(obj, int | float | complex | bool | list): +# if copy is False: +# raise ValueError( +# "copy=False isn't supported for scalar inputs and Python lists" +# ) +# obj = np.asarray(obj) +# if fill_value is None: +# fill_value = 0.0 + +# if _is_scipy_sparse_obj(obj): # scipy constructor +# jl_data = self._from_scipy_sparse(obj, fill_value=fill_value, copy=copy) +# self._obj = jl_data +# elif isinstance(obj, np.ndarray): # numpy constructor +# jl_data = self._from_numpy(obj, fill_value=fill_value, copy=copy) +# self._obj = jl_data +# elif isinstance(obj, Storage): # from-storage constructor +# if copy: +# self._raise_julia_copy_not_supported() +# order = self.preprocess_order( +# obj.order, self.get_lvl_ndim(obj.levels_descr._obj) +# ) +# self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order) +# elif jl.isa(obj, jl.Finch.Tensor): # raw-Julia-object constructors +# if copy: +# self._raise_julia_copy_not_supported() +# self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1))) +# elif jl.isa(obj, jl.Finch.Tensor): +# if copy: +# self._raise_julia_copy_not_supported() +# self._obj = obj +# elif isinstance(obj, FinchJLTensor): +# self._obj = obj._obj +# else: +# raise ValueError( +# "Either scalar, numpy, scipy.sparse or a raw julia object should " +# f"be provided. Found: {type(obj)}" +# ) + +# @property +# def element_type(self): +# return jl.eltype(self._obj.body) + +# @property +# def dtype(self) -> DType: +# return jl.eltype(self._obj.body) + +# @property +# def ndim(self) -> int: +# return jl.ndims(self._obj) + +# @property +# def shape(self) -> tuple[int, ...]: +# return jl.size(self._obj) + +# @property +# def size(self) -> int: +# return np.prod(self.shape) + +# @property +# def fill_value(self) -> np.number: +# return jl.fill_value(self._obj) + +# @property +# def _is_dense(self) -> bool: +# lvl = self._obj.body.lvl +# for _ in self.shape: +# if not jl.isa(lvl, jl.Finch.Dense): +# return False +# lvl = lvl.lvl +# return True + +# @property +# def _order(self) -> tuple[int, ...]: +# return jl.typeof(self._obj).parameters[1] + +# @property +# def mT(self) -> Tensor: +# axes = list(range(self.ndim)) +# axes[-2], axes[-1] = axes[-1], axes[-2] +# axes = tuple(axes) +# return self.permute_dims(axes) + +# @property +# def device(self) -> str: +# return "cpu" + +# def to_device( +# self, device: Device, /, *, stream: int | Any | None = None +# ) -> Tensor: +# if device != "cpu": +# raise ValueError("Only `device='cpu'` is supported.") + +# return self + +# @classmethod +# def get_lvl_ndim(cls, lvl: JuliaObj) -> int: +# ndim = 0 +# while True: +# ndim += 1 +# lvl = lvl.lvl +# if jl.isa(lvl, jl.Finch.Element): +# break +# return ndim + +# def todense(self) -> np.ndarray: +# obj = self._obj + +# if self._is_dense: +# # don't materialize a dense finch tensor +# shape = jl.size(obj.body) +# dense_tensor = obj.body.lvl +# else: +# # create materialized dense array +# shape = jl.size(obj) +# dense_lvls = jl.Element(jc.convert(self.dtype, jl.fill_value(obj))) +# for _ in range(self.ndim): +# dense_lvls = jl.Dense(dense_lvls) +# dense_tensor = jl.Tensor(dense_lvls, obj).lvl # materialize + +# for _ in range(self.ndim): +# dense_tensor = dense_tensor.lvl + +# result = np.asarray(jl.reshape(dense_tensor.val, shape)) +# return result.transpose(self.get_order()) if self._is_dense else result + +# #TODO: Do we need? +# def permute_dims(self, axes: tuple[int, ...]) -> Tensor: +# axes = tuple(i + 1 for i in axes) +# new_obj = jl.permutedims(self._obj, axes) +# return Tensor(new_obj) + +# def to_storage(self, storage: Storage) -> Tensor: +# return Tensor(self._from_other_tensor(self, storage=storage)) + +# @classmethod +# def _from_other_tensor(cls, tensor: Tensor, storage: Storage) -> JuliaObj: +# order = cls.preprocess_order(storage.order, tensor.ndim) +# result = jl.copyto_b( +# jl.swizzle(jl.Tensor(storage.levels_descr._obj), *order), tensor._obj +# ) +# return jl.dropfills(result) if tensor._is_dense else result + +# @classmethod +# def _from_numpy( +# cls, arr: np.ndarray, fill_value: np.number, copy: bool | None = None +# ) -> JuliaObj: +# if copy: +# arr = arr.copy() +# order_char = "F" if np.isfortran(arr) else "C" +# order = cls.preprocess_order(order_char, arr.ndim) +# inv_order = tuple(i - 1 for i in jl.invperm(order)) + +# dtype = arr.dtype.type +# if ( +# dtype == np.bool_ +# ): # Fails with: Finch currently only supports isbits defaults +# dtype = jl_dtypes.bool +# fill_value = dtype(fill_value) +# lvl = Element(fill_value, arr.reshape(-1, order=order_char)) +# for i in inv_order: +# lvl = Dense(lvl, arr.shape[i]) +# return jl.swizzle(jl.Tensor(lvl._obj), *order) + +# @classmethod +# def from_scipy_sparse( +# cls, +# x, +# fill_value: np.number | None = None, +# copy: bool | None = None, +# ) -> Tensor: +# if not _is_scipy_sparse_obj(x): +# raise ValueError("{x} is not a SciPy sparse object.") +# return Tensor(x, fill_value=fill_value, copy=copy) + +# @classmethod +# def _from_scipy_sparse( +# cls, +# x, +# *, +# fill_value: np.number | None = None, +# copy: bool | None = None, +# ) -> JuliaObj: +# if copy is False and not ( +# x.format in ("coo", "csr", "csc") and x.has_canonical_format +# ): +# raise ValueError( +# "Unable to avoid copy while creating an array as requested." +# ) +# if x.format not in ("coo", "csr", "csc"): +# x = x.asformat("coo") +# if copy: +# x = x.copy() +# if not x.has_canonical_format: +# x.sum_duplicates() +# assert x.has_canonical_format + +# if x.format == "coo": +# return cls.construct_coo_jl_object( +# coords=(x.col, x.row), +# data=x.data, +# shape=x.shape[::-1], +# order=Tensor.row_major, +# fill_value=fill_value, +# ) +# if x.format == "csc": +# return cls.construct_csc_jl_object( +# arg=(x.data, x.indices, x.indptr), +# shape=x.shape, +# fill_value=fill_value, +# ) +# if x.format == "csr": +# return cls.construct_csr_jl_object( +# arg=(x.data, x.indices, x.indptr), +# shape=x.shape, +# fill_value=fill_value, +# ) +# raise ValueError(f"Unsupported SciPy format: {type(x)}") + +# @classmethod +# def construct_coo_jl_object( +# cls, coords, data, shape, order, fill_value=0.0 +# ) -> JuliaObj: +# assert len(coords) == 2 +# ndim = len(shape) +# order = cls.preprocess_order(order, ndim) + +# lvl = jl.Element(data.dtype.type(fill_value), data) +# ptr = jl.Vector[jl.Int]([1, len(data) + 1]) +# tbl = tuple(jl.PlusOneVector(arr) for arr in coords) + +# return jl.swizzle(jl.Tensor(jl.SparseCOO[ndim](lvl, shape, ptr, tbl)), *order) + +# @classmethod +# def construct_coo( +# cls, coords, data, shape, order=row_major, fill_value=0.0 +# ) -> Tensor: +# return Tensor( +# cls.construct_coo_jl_object(coords, data, shape, order, fill_value) +# ) + +# @staticmethod +# def _construct_compressed2d_jl_object( +# arg: TupleOf3Arrays, +# shape: tuple[int, ...], +# order: tuple[int, ...], +# fill_value: np.number = 0.0, +# ) -> JuliaObj: +# assert isinstance(arg, tuple) and len(arg) == 3 +# assert len(shape) == 2 + +# data, indices, indptr = arg +# dtype = data.dtype.type +# indices = jl.PlusOneVector(indices) +# indptr = jl.PlusOneVector(indptr) + +# lvl = jl.Element(dtype(fill_value), data) +# return jl.swizzle( +# jl.Tensor( +# jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1]) +# ), +# *order, +# ) + +# @classmethod +# def construct_csc_jl_object( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> JuliaObj: +# return cls._construct_compressed2d_jl_object( +# arg=arg, shape=shape, order=(1, 2), fill_value=fill_value +# ) + +# @classmethod +# def construct_csc( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> Tensor: +# return Tensor(cls.construct_csc_jl_object(arg, shape, fill_value)) + +# @classmethod +# def construct_csr_jl_object( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> JuliaObj: +# return cls._construct_compressed2d_jl_object( +# arg=arg, shape=shape[::-1], order=(2, 1), fill_value=fill_value +# ) + +# @classmethod +# def construct_csr( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> Tensor: +# return Tensor(cls.construct_csr_jl_object(arg, shape, fill_value)) + +# @staticmethod +# def construct_csf_jl_object( +# arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> JuliaObj: +# assert isinstance(arg, tuple) and len(arg) == 3 + +# data, indices_list, indptr_list = arg +# dtype = data.dtype.type + +# assert len(indices_list) == len(shape) - 1 +# assert len(indptr_list) == len(shape) - 1 + +# indices_list = [jl.PlusOneVector(i) for i in indices_list] +# indptr_list = [jl.PlusOneVector(i) for i in indptr_list] + +# lvl = jl.Element(dtype(fill_value), data) +# for size, indices, indptr in zip( +# shape[:-1], indices_list, indptr_list, strict=False +# ): +# lvl = jl.SparseList(lvl, size, indptr, indices) + +# return jl.swizzle( +# jl.Tensor(jl.Dense(lvl, shape[-1])), *range(1, len(shape) + 1) +# ) + +# @classmethod +# def construct_csf( +# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 +# ) -> Tensor: +# return Tensor(cls.construct_csf_jl_object(arg, shape, fill_value)) + +# def to_scipy_sparse(self, accept_fv=None): +# import scipy.sparse as sp + +# if accept_fv is None: +# accept_fv = [0] +# elif not isinstance(accept_fv, Iterable): +# accept_fv = [accept_fv] + +# if self.ndim != 2: +# raise ValueError( +# "Can only convert a 2-dimensional array to a Scipy sparse matrix." +# ) +# if not builtins.any(_eq_scalars(self.fill_value, fv) for fv in accept_fv): +# raise ValueError( +# f"Can only convert arrays with {accept_fv} fill-values " +# "to a Scipy sparse matrix." +# ) +# order = self.get_order() +# body = self._obj.body + +# if str(jl.typeof(body.lvl).name.name) == "SparseCOOLevel": +# data = np.asarray(body.lvl.lvl.val) +# coords = body.lvl.tbl +# row, col = coords[::-1] if order == (1, 0) else coords +# row, col = np.asarray(row) - 1, np.asarray(col) - 1 +# return sp.coo_matrix((data, (row, col)), shape=self.shape) + +# if ( +# str(jl.typeof(body.lvl).name.name) == "DenseLevel" +# and str(jl.typeof(body.lvl.lvl).name.name) == "SparseListLevel" +# ): +# data = np.asarray(body.lvl.lvl.lvl.val) +# indices = np.asarray(body.lvl.lvl.idx) - 1 +# indptr = np.asarray(body.lvl.lvl.ptr) - 1 +# sp_class = sp.csr_matrix if order == (1, 0) else sp.csc_matrix +# return sp_class((data, indices, indptr), shape=self.shape) +# if ( +# jl.typeof(body.lvl).name.name in sparse_formats_names +# or jl.typeof(body.lvl.lvl).name.name in sparse_formats_names +# ): +# storage = Storage(SparseCOO(self.ndim, Element(self.fill_value)), order) +# return self.to_storage(storage).to_scipy_sparse() +# raise ValueError("Tensor can't be converted to scipy.sparse object.") + +# @staticmethod +# def _raise_julia_copy_not_supported() -> None: +# raise ValueError("copy=True isn't supported for Julia object inputs") + +# def __array_namespace__(self, *, api_version: str | None = None) -> Any: +# if api_version is None: +# api_version = "2024.12" + +# if api_version not in {"2021.12", "2022.12", "2023.12", "2024.12"}: +# raise ValueError(f'"{api_version}" Array API version not supported.') +# import finch + +# return finch + + +# def random(shape, density=0.01, random_state=None): +# args = [*shape, density] +# if random_state is not None: +# if isinstance(random_state, np.random.Generator): +# seed = random_state.integers(np.iinfo(np.int32).max) +# else: +# seed = random_state +# rng = jl.Random.default_rng() +# jl.Random.seed_b(rng, seed) +# args = [rng] + args +# return Tensor(jl.fsprand(*args)) + + +# def asarray( +# obj, +# /, +# *, +# dtype: DType | None = None, +# format: str | None = None, +# fill_value: np.number | None = None, +# device: Device | None = None, +# copy: bool | None = None, +# ) -> Tensor: +# if format not in {"coo", "csr", "csc", "csf", "dense", None}: +# raise ValueError(f"{format} format not supported.") +# _validate_device(device) +# tensor = ( +# obj +# if isinstance(obj, Tensor) +# else Tensor(obj, fill_value=fill_value, copy=copy) +# ) +# if format is not None: +# if copy is False: +# raise ValueError( +# "Unable to avoid copy while creating an array as requested." +# ) +# order = tensor.get_order() +# if format == "coo": +# storage = Storage(SparseCOO(tensor.ndim, Element(tensor.fill_value)), order) +# elif format == "csr": +# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (2, 1)) +# elif format == "csc": +# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (1, 2)) +# elif format == "csf": +# storage = Element(tensor.fill_value) +# for _ in range(tensor.ndim - 1): +# storage = SparseList(storage) +# storage = Storage(Dense(storage), order) +# elif format == "dense": +# storage = DenseStorage(tensor.ndim, tensor.dtype, order) +# tensor = tensor.to_storage(storage) + +# if dtype is not None: +# return astype(tensor, dtype, copy=copy) +# return tensor + + +# def reshape( +# x: Tensor, /, shape: tuple[int, ...], *, copy: bool | None = None +# ) -> Tensor: +# if copy is False: +# raise ValueError("Unable to avoid copy during reshape.") +# # TODO: https://github.com/finch-tensor/Finch.jl/issues/743 +# # Revert to `jl.reshape` implementation once aforementioned +# # issue is solved. +# warnings.warn( +# "`reshape` densified the input tensor.", PerformanceWarning, stacklevel=2 +# ) +# arr = x.todense() +# arr = arr.reshape(shape) +# return Tensor(arr) + + +# def _validate_device(device: Device) -> None: +# if device not in {"cpu", None}: +# raise ValueError( +# f'Device not understood. Only "cpu" is allowed, but received: {device}' +# ) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index afdffb1..d7a094d 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,577 +1,117 @@ -# from __future__ import annotations - -# import builtins -# import warnings -# from collections.abc import Callable, Iterable -# from typing import Any, Literal - -# import numpy as np -# from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple - -# from . import dtypes as jl_dtypes -# from .errors import PerformanceWarning -# from .julia import jc, jl -# from .levels import ( -# Dense, -# DenseStorage, -# Element, -# SparseCOO, -# SparseList, -# Storage, -# _Display, -# sparse_formats_names, -# ) -# from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix -# from finchlite import Tensor, TensorFType, EagerTensor - - -# class FinchJLTensorFType(TensorFType): -# def __init__(self, jltype): -# # Julia type associated with the tensor -# self.jltype = jltype - -# def ndims(self) -> np.intp: -# return np.intp(jl.ndims(self.jltype)) +import numpy as np +from typing import Any +from finchlite import EagerTensor, TensorFType, Tensor + +from .julia import jc, jl +from .levels import ( + Dense, + DenseStorage, + Element, + SparseCOO, + SparseList, + Storage, + _Display, + sparse_formats_names, +) + +class FinchJLTensorFType(TensorFType): + def __init__(self, jltype, shape_type): + # Julia type associated with the tensor + self.jltype = jltype + self._shape_type = shape_type + + def ndims(self) -> np.intp: + return np.intp(jl.ndims(self.jltype)) -# def fill_value(self) -> Any: -# return jl.fill_value(self.jltype) + def fill_value(self) -> Any: + return jl.fill_value(self.jltype) -# def element_type(self) -> Any: -# return jl.eltype(self.jltype) + def element_type(self) -> Any: + return jl.eltype(self.jltype) -# # TODO: implement later -# def shape_type(self) -> tuple[type, ...]: -# ... - -# def __call__(self, shape: tuple) -> Tensor: -# ... - -# def from_numpy(self, arr: np.ndarray) -> Tensor: -# ... - -# class FinchJLTensor(_Display, EagerTensor): -# """ -# A wrapper class for Finch.Tensor and Finch.SwizzleArray. - -# Constructors -# ------------ -# FinchJLTensor(scipy.sparse.spmatrix) -# Construct a Tensor out of a `scipy.sparse` object. Supported formats are: `COO`, -# `CSC`, and `CSR`. -# FinchJLTensor(numpy.ndarray) -# Construct a Tensor out of a NumPy array object. This is a no-copy operation. -# FinchJLTensor(Storage) -# Initialize a Tensor with a `storage` description. `storage` can already hold -# data. -# FinchJLTensor(julia_object) -# Tensor created from a compatible raw Julia object. Must be a `Tensor`. -# This is a no-copy operation. - -# Parameters -# ---------- -# obj : np.ndarray or scipy.sparse or Storage or Finch.Tensor -# Input to construct a Tensor. It's a no-copy operation of for NumPy and -# SciPy input. For Storage it's levels' description with order. The order -# numbers the dimensions from the fastest to slowest. The leaf nodes have -# mode `0` and the root node has mode `n-1`. If the tensor was square of -# size `N`, then `N .^ order == strides`. Available options are "C" -# (row-major), "F" (column-major), or a custom order. Default: row-major. -# fill_value : np.number, optional -# Only used when `numpy.ndarray` or `scipy.sparse` is passed. -# copy : bool, optional -# If ``True``, then the object is copied. If ``None`` then the object is -# copied only if needed. For ``False`` it raises a ``ValueError`` if a -# copy cannot be avoided. Default: ``None``. - -# Returns -# ------- -# FinchJLTensor -# Python wrapper for Finch.jl `Tensor`. - -# Examples -# -------- -# >>> import numpy as np -# >>> import finch -# >>> arr2d = np.arange(6).reshape((2, 3)) -# >>> t1 = finch.FinchJLTensor(arr2d) -# >>> t1.todense() -# array([[0, 1, 2], -# [3, 4, 5]]) -# >>> np.shares_memory(t1.todense(), arr2d) -# True -# >>> storage = finch.Storage( -# ... finch.Dense(finch.SparseList(finch.Element(1))), order="C" -# ... ) -# >>> t2 = t1.to_storage(storage) -# >>> t2.todense() -# array([[0, 1, 2], -# [3, 4, 5]]) -# """ - -# def __init__( -# self, -# obj: np.ndarray | spmatrix | Storage | JuliaObj, -# /, -# *, -# fill_value: np.number | None = None, -# copy: bool | None = None, -# ): -# if isinstance(obj, int | float | complex | bool | list): -# if copy is False: -# raise ValueError( -# "copy=False isn't supported for scalar inputs and Python lists" -# ) -# obj = np.asarray(obj) -# if fill_value is None: -# fill_value = 0.0 - -# if _is_scipy_sparse_obj(obj): # scipy constructor -# jl_data = self._from_scipy_sparse(obj, fill_value=fill_value, copy=copy) -# self._obj = jl_data -# elif isinstance(obj, np.ndarray): # numpy constructor -# jl_data = self._from_numpy(obj, fill_value=fill_value, copy=copy) -# self._obj = jl_data -# elif isinstance(obj, Storage): # from-storage constructor -# if copy: -# self._raise_julia_copy_not_supported() -# order = self.preprocess_order( -# obj.order, self.get_lvl_ndim(obj.levels_descr._obj) -# ) -# self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order) -# elif jl.isa(obj, jl.Finch.Tensor): # raw-Julia-object constructors -# if copy: -# self._raise_julia_copy_not_supported() -# self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1))) -# elif jl.isa(obj, jl.Finch.Tensor): -# if copy: -# self._raise_julia_copy_not_supported() -# self._obj = obj -# elif isinstance(obj, FinchJLTensor): -# self._obj = obj._obj -# else: -# raise ValueError( -# "Either scalar, numpy, scipy.sparse or a raw julia object should " -# f"be provided. Found: {type(obj)}" -# ) - -# @property -# def element_type(self): -# return jl.eltype(self._obj.body) - -# @property -# def dtype(self) -> DType: -# return jl.eltype(self._obj.body) - -# @property -# def ndim(self) -> int: -# return jl.ndims(self._obj) - -# @property -# def shape(self) -> tuple[int, ...]: -# return jl.size(self._obj) - -# @property -# def size(self) -> int: -# return np.prod(self.shape) - -# @property -# def fill_value(self) -> np.number: -# return jl.fill_value(self._obj) - -# @property -# def _is_dense(self) -> bool: -# lvl = self._obj.body.lvl -# for _ in self.shape: -# if not jl.isa(lvl, jl.Finch.Dense): -# return False -# lvl = lvl.lvl -# return True - -# @property -# def _order(self) -> tuple[int, ...]: -# return jl.typeof(self._obj).parameters[1] - -# @property -# def mT(self) -> Tensor: -# axes = list(range(self.ndim)) -# axes[-2], axes[-1] = axes[-1], axes[-2] -# axes = tuple(axes) -# return self.permute_dims(axes) - -# @property -# def device(self) -> str: -# return "cpu" - -# def to_device( -# self, device: Device, /, *, stream: int | Any | None = None -# ) -> Tensor: -# if device != "cpu": -# raise ValueError("Only `device='cpu'` is supported.") - -# return self - -# @classmethod -# def get_lvl_ndim(cls, lvl: JuliaObj) -> int: -# ndim = 0 -# while True: -# ndim += 1 -# lvl = lvl.lvl -# if jl.isa(lvl, jl.Finch.Element): -# break -# return ndim - -# def todense(self) -> np.ndarray: -# obj = self._obj - -# if self._is_dense: -# # don't materialize a dense finch tensor -# shape = jl.size(obj.body) -# dense_tensor = obj.body.lvl -# else: -# # create materialized dense array -# shape = jl.size(obj) -# dense_lvls = jl.Element(jc.convert(self.dtype, jl.fill_value(obj))) -# for _ in range(self.ndim): -# dense_lvls = jl.Dense(dense_lvls) -# dense_tensor = jl.Tensor(dense_lvls, obj).lvl # materialize - -# for _ in range(self.ndim): -# dense_tensor = dense_tensor.lvl - -# result = np.asarray(jl.reshape(dense_tensor.val, shape)) -# return result.transpose(self.get_order()) if self._is_dense else result - -# #TODO: Do we need? -# def permute_dims(self, axes: tuple[int, ...]) -> Tensor: -# axes = tuple(i + 1 for i in axes) -# new_obj = jl.permutedims(self._obj, axes) -# return Tensor(new_obj) - -# def to_storage(self, storage: Storage) -> Tensor: -# return Tensor(self._from_other_tensor(self, storage=storage)) - -# @classmethod -# def _from_other_tensor(cls, tensor: Tensor, storage: Storage) -> JuliaObj: -# order = cls.preprocess_order(storage.order, tensor.ndim) -# result = jl.copyto_b( -# jl.swizzle(jl.Tensor(storage.levels_descr._obj), *order), tensor._obj -# ) -# return jl.dropfills(result) if tensor._is_dense else result - -# @classmethod -# def _from_numpy( -# cls, arr: np.ndarray, fill_value: np.number, copy: bool | None = None -# ) -> JuliaObj: -# if copy: -# arr = arr.copy() -# order_char = "F" if np.isfortran(arr) else "C" -# order = cls.preprocess_order(order_char, arr.ndim) -# inv_order = tuple(i - 1 for i in jl.invperm(order)) - -# dtype = arr.dtype.type -# if ( -# dtype == np.bool_ -# ): # Fails with: Finch currently only supports isbits defaults -# dtype = jl_dtypes.bool -# fill_value = dtype(fill_value) -# lvl = Element(fill_value, arr.reshape(-1, order=order_char)) -# for i in inv_order: -# lvl = Dense(lvl, arr.shape[i]) -# return jl.swizzle(jl.Tensor(lvl._obj), *order) - -# @classmethod -# def from_scipy_sparse( -# cls, -# x, -# fill_value: np.number | None = None, -# copy: bool | None = None, -# ) -> Tensor: -# if not _is_scipy_sparse_obj(x): -# raise ValueError("{x} is not a SciPy sparse object.") -# return Tensor(x, fill_value=fill_value, copy=copy) - -# @classmethod -# def _from_scipy_sparse( -# cls, -# x, -# *, -# fill_value: np.number | None = None, -# copy: bool | None = None, -# ) -> JuliaObj: -# if copy is False and not ( -# x.format in ("coo", "csr", "csc") and x.has_canonical_format -# ): -# raise ValueError( -# "Unable to avoid copy while creating an array as requested." -# ) -# if x.format not in ("coo", "csr", "csc"): -# x = x.asformat("coo") -# if copy: -# x = x.copy() -# if not x.has_canonical_format: -# x.sum_duplicates() -# assert x.has_canonical_format - -# if x.format == "coo": -# return cls.construct_coo_jl_object( -# coords=(x.col, x.row), -# data=x.data, -# shape=x.shape[::-1], -# order=Tensor.row_major, -# fill_value=fill_value, -# ) -# if x.format == "csc": -# return cls.construct_csc_jl_object( -# arg=(x.data, x.indices, x.indptr), -# shape=x.shape, -# fill_value=fill_value, -# ) -# if x.format == "csr": -# return cls.construct_csr_jl_object( -# arg=(x.data, x.indices, x.indptr), -# shape=x.shape, -# fill_value=fill_value, -# ) -# raise ValueError(f"Unsupported SciPy format: {type(x)}") - -# @classmethod -# def construct_coo_jl_object( -# cls, coords, data, shape, order, fill_value=0.0 -# ) -> JuliaObj: -# assert len(coords) == 2 -# ndim = len(shape) -# order = cls.preprocess_order(order, ndim) - -# lvl = jl.Element(data.dtype.type(fill_value), data) -# ptr = jl.Vector[jl.Int]([1, len(data) + 1]) -# tbl = tuple(jl.PlusOneVector(arr) for arr in coords) - -# return jl.swizzle(jl.Tensor(jl.SparseCOO[ndim](lvl, shape, ptr, tbl)), *order) - -# @classmethod -# def construct_coo( -# cls, coords, data, shape, order=row_major, fill_value=0.0 -# ) -> Tensor: -# return Tensor( -# cls.construct_coo_jl_object(coords, data, shape, order, fill_value) -# ) - -# @staticmethod -# def _construct_compressed2d_jl_object( -# arg: TupleOf3Arrays, -# shape: tuple[int, ...], -# order: tuple[int, ...], -# fill_value: np.number = 0.0, -# ) -> JuliaObj: -# assert isinstance(arg, tuple) and len(arg) == 3 -# assert len(shape) == 2 - -# data, indices, indptr = arg -# dtype = data.dtype.type -# indices = jl.PlusOneVector(indices) -# indptr = jl.PlusOneVector(indptr) - -# lvl = jl.Element(dtype(fill_value), data) -# return jl.swizzle( -# jl.Tensor( -# jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1]) -# ), -# *order, -# ) - -# @classmethod -# def construct_csc_jl_object( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> JuliaObj: -# return cls._construct_compressed2d_jl_object( -# arg=arg, shape=shape, order=(1, 2), fill_value=fill_value -# ) - -# @classmethod -# def construct_csc( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> Tensor: -# return Tensor(cls.construct_csc_jl_object(arg, shape, fill_value)) - -# @classmethod -# def construct_csr_jl_object( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> JuliaObj: -# return cls._construct_compressed2d_jl_object( -# arg=arg, shape=shape[::-1], order=(2, 1), fill_value=fill_value -# ) - -# @classmethod -# def construct_csr( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> Tensor: -# return Tensor(cls.construct_csr_jl_object(arg, shape, fill_value)) - -# @staticmethod -# def construct_csf_jl_object( -# arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> JuliaObj: -# assert isinstance(arg, tuple) and len(arg) == 3 - -# data, indices_list, indptr_list = arg -# dtype = data.dtype.type - -# assert len(indices_list) == len(shape) - 1 -# assert len(indptr_list) == len(shape) - 1 - -# indices_list = [jl.PlusOneVector(i) for i in indices_list] -# indptr_list = [jl.PlusOneVector(i) for i in indptr_list] - -# lvl = jl.Element(dtype(fill_value), data) -# for size, indices, indptr in zip( -# shape[:-1], indices_list, indptr_list, strict=False -# ): -# lvl = jl.SparseList(lvl, size, indptr, indices) - -# return jl.swizzle( -# jl.Tensor(jl.Dense(lvl, shape[-1])), *range(1, len(shape) + 1) -# ) - -# @classmethod -# def construct_csf( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> Tensor: -# return Tensor(cls.construct_csf_jl_object(arg, shape, fill_value)) - -# def to_scipy_sparse(self, accept_fv=None): -# import scipy.sparse as sp - -# if accept_fv is None: -# accept_fv = [0] -# elif not isinstance(accept_fv, Iterable): -# accept_fv = [accept_fv] - -# if self.ndim != 2: -# raise ValueError( -# "Can only convert a 2-dimensional array to a Scipy sparse matrix." -# ) -# if not builtins.any(_eq_scalars(self.fill_value, fv) for fv in accept_fv): -# raise ValueError( -# f"Can only convert arrays with {accept_fv} fill-values " -# "to a Scipy sparse matrix." -# ) -# order = self.get_order() -# body = self._obj.body - -# if str(jl.typeof(body.lvl).name.name) == "SparseCOOLevel": -# data = np.asarray(body.lvl.lvl.val) -# coords = body.lvl.tbl -# row, col = coords[::-1] if order == (1, 0) else coords -# row, col = np.asarray(row) - 1, np.asarray(col) - 1 -# return sp.coo_matrix((data, (row, col)), shape=self.shape) - -# if ( -# str(jl.typeof(body.lvl).name.name) == "DenseLevel" -# and str(jl.typeof(body.lvl.lvl).name.name) == "SparseListLevel" -# ): -# data = np.asarray(body.lvl.lvl.lvl.val) -# indices = np.asarray(body.lvl.lvl.idx) - 1 -# indptr = np.asarray(body.lvl.lvl.ptr) - 1 -# sp_class = sp.csr_matrix if order == (1, 0) else sp.csc_matrix -# return sp_class((data, indices, indptr), shape=self.shape) -# if ( -# jl.typeof(body.lvl).name.name in sparse_formats_names -# or jl.typeof(body.lvl.lvl).name.name in sparse_formats_names -# ): -# storage = Storage(SparseCOO(self.ndim, Element(self.fill_value)), order) -# return self.to_storage(storage).to_scipy_sparse() -# raise ValueError("Tensor can't be converted to scipy.sparse object.") - -# @staticmethod -# def _raise_julia_copy_not_supported() -> None: -# raise ValueError("copy=True isn't supported for Julia object inputs") - -# def __array_namespace__(self, *, api_version: str | None = None) -> Any: -# if api_version is None: -# api_version = "2024.12" - -# if api_version not in {"2021.12", "2022.12", "2023.12", "2024.12"}: -# raise ValueError(f'"{api_version}" Array API version not supported.') -# import finch - -# return finch - - -# def random(shape, density=0.01, random_state=None): -# args = [*shape, density] -# if random_state is not None: -# if isinstance(random_state, np.random.Generator): -# seed = random_state.integers(np.iinfo(np.int32).max) -# else: -# seed = random_state -# rng = jl.Random.default_rng() -# jl.Random.seed_b(rng, seed) -# args = [rng] + args -# return Tensor(jl.fsprand(*args)) - - -# def asarray( -# obj, -# /, -# *, -# dtype: DType | None = None, -# format: str | None = None, -# fill_value: np.number | None = None, -# device: Device | None = None, -# copy: bool | None = None, -# ) -> Tensor: -# if format not in {"coo", "csr", "csc", "csf", "dense", None}: -# raise ValueError(f"{format} format not supported.") -# _validate_device(device) -# tensor = ( -# obj -# if isinstance(obj, Tensor) -# else Tensor(obj, fill_value=fill_value, copy=copy) -# ) -# if format is not None: -# if copy is False: -# raise ValueError( -# "Unable to avoid copy while creating an array as requested." -# ) -# order = tensor.get_order() -# if format == "coo": -# storage = Storage(SparseCOO(tensor.ndim, Element(tensor.fill_value)), order) -# elif format == "csr": -# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (2, 1)) -# elif format == "csc": -# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (1, 2)) -# elif format == "csf": -# storage = Element(tensor.fill_value) -# for _ in range(tensor.ndim - 1): -# storage = SparseList(storage) -# storage = Storage(Dense(storage), order) -# elif format == "dense": -# storage = DenseStorage(tensor.ndim, tensor.dtype, order) -# tensor = tensor.to_storage(storage) - -# if dtype is not None: -# return astype(tensor, dtype, copy=copy) -# return tensor - - -# def reshape( -# x: Tensor, /, shape: tuple[int, ...], *, copy: bool | None = None -# ) -> Tensor: -# if copy is False: -# raise ValueError("Unable to avoid copy during reshape.") -# # TODO: https://github.com/finch-tensor/Finch.jl/issues/743 -# # Revert to `jl.reshape` implementation once aforementioned -# # issue is solved. -# warnings.warn( -# "`reshape` densified the input tensor.", PerformanceWarning, stacklevel=2 -# ) -# arr = x.todense() -# arr = arr.reshape(shape) -# return Tensor(arr) - - -# def _validate_device(device: Device) -> None: -# if device not in {"cpu", None}: -# raise ValueError( -# f'Device not understood. Only "cpu" is allowed, but received: {device}' -# ) + def shape_type(self) -> tuple[type, ...]: + return self._shape_type + + def __call__(self, shape: tuple) -> Tensor: + return FinchJLTensor(np.ones(shape=shape)) + + def from_numpy(self, arr: np.ndarray) -> Tensor: + return FinchJLTensor(arr) + +# TODO: Do we need the scipy and raw julia stuff +class FinchJLTensor(_Display, EagerTensor): + def __init__( + self, + obj: np.ndarray, + /, + *, + fill_value: np.number | None = None, + copy: bool | None = None, + ): + if isinstance(obj, int | float | complex | bool | list): + if copy is False: + raise ValueError( + "copy=False isn't supported for scalar inputs and Python lists" + ) + obj = np.asarray(obj) + if fill_value is None: + fill_value = 0.0 + + if isinstance(obj, np.ndarray): # numpy constructor + jl_data = self._from_numpy(obj, fill_value=fill_value, copy=copy) + self._shape = obj.shape + self._obj = jl_data + else: + raise ValueError( + "Either scalar, numpy, scipy.sparse or a raw julia object should " + f"be provided. Found: {type(obj)}" + ) + + @property + def ftype(self): + """ + Returns the ftype of the buffer, which is a BufferizedNDArrayFType. + """ + shape_type = [] + for idx in self._shape: + shape_type.append(type(idx)) + return FinchJLTensorFType(jltype=jl.typeof(self._obj), shape_type=shape_type) + + @property + def shape(self) -> tuple: + """Shape of the tensor.""" + return self._shape + + # TODO: do we need to have all the order stuff still? + @classmethod + def _from_numpy( + cls, arr: np.ndarray, fill_value: np.number, copy: bool | None = None + ) -> JuliaObj: + if copy: + arr = arr.copy() + order_char = "F" if np.isfortran(arr) else "C" + order = cls.preprocess_order(order_char, arr.ndim) + inv_order = tuple(i - 1 for i in jl.invperm(order)) + + dtype = arr.dtype.type + if ( + dtype == np.bool_ + ): # Fails with: Finch currently only supports isbits defaults + dtype = jl_dtypes.bool + fill_value = dtype(fill_value) + lvl = Element(fill_value, arr.reshape(-1, order=order_char)) + for i in inv_order: + lvl = Dense(lvl, arr.shape[i]) + return jl.swizzle(jl.Tensor(lvl._obj), *order) + + @classmethod + def preprocess_order(cls, order: str, ndim: int) -> tuple[int, ...]: + if order == 'F': + permutation = tuple(range(1, ndim + 1)) + elif order == 'C': + permutation = tuple(range(1, ndim + 1)[::-1]) + else: + raise ValueError( + f"order must be 'C' or 'F'." + ) + return permutation From 157cf03a8ca447b3a673010d185aaefd39001743 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 11 Feb 2026 12:21:26 -0500 Subject: [PATCH 07/81] feat: FinchJLTensor works! --- src/finch/compiler.py | 15 ++++++++------- src/finch/levels.py | 2 +- src/finch/tensor.py | 11 ++++++++++- src/finch/typing.py | 2 ++ tests/test_compiler.py | 5 +++-- 5 files changed, 24 insertions(+), 11 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index a47835a..08f8d2e 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,3 +1,5 @@ +from finch.tensor import FinchJLTensor + from finchlite.compile import NotationCompiler from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary import finchlite.finch_notation.nodes as ntn @@ -13,14 +15,13 @@ class FinchJLKernel(AssemblyKernel): def __init__(self, func_name, jl_code): + # We store this code so that we can verify it in pytest self.jl_code = jl_code self.func_name = func_name - - print(jl_code) - # jl.seval(jl_code) + jl.seval(self.jl_code) # TODO: Switch back to (self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...] - def __call__(self, *args: tuple[Any, ...]): + def __call__(self, *args: tuple[FinchJLTensor, ...]): argList = [] for arg in args: argList.append(f"arg{len(argList)}") @@ -70,7 +71,7 @@ def generate_julia(self, prgm, nestingLvl=0): # finding loop bounds if isinstance(rhs, ntn.Call) and rhs.op.val == dimension: return "" - + tab_str = " " * nestingLvl return f"{tab_str}{self.generate_julia(lhs, nestingLvl)} = {self.generate_julia(rhs, nestingLvl)}" @@ -85,7 +86,7 @@ def generate_julia(self, prgm, nestingLvl=0): case ntn.Loop(idx, _, body): tab_str = " " * nestingLvl - tab_str_1 = " " * (nestingLvl+1) + tab_str_1 = " " * (nestingLvl + 1) is_outermost_loop = False if self.in_finch_block is False: @@ -94,7 +95,7 @@ def generate_julia(self, prgm, nestingLvl=0): loop_body = self.generate_julia(body, nestingLvl + 2) else: loop_body = self.generate_julia(body, nestingLvl + 1) - + if not is_outermost_loop: return f"{tab_str}for {idx.name} = _\n{loop_body}{tab_str}end\n" else: diff --git a/src/finch/levels.py b/src/finch/levels.py index eb10eaa..07d6142 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -1,5 +1,5 @@ from .julia import jl -from .typing import DType, JuliaObj +from .typing import DType, JuliaObj, OrderType class _Display: diff --git a/src/finch/tensor.py b/src/finch/tensor.py index d7a094d..335818f 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -2,6 +2,7 @@ from typing import Any from finchlite import EagerTensor, TensorFType, Tensor +from .typing import OrderType, JuliaObj from .julia import jc, jl from .levels import ( Dense, @@ -37,6 +38,14 @@ def __call__(self, shape: tuple) -> Tensor: def from_numpy(self, arr: np.ndarray) -> Tensor: return FinchJLTensor(arr) + + def __eq__(self, other): + if not isinstance(other, FinchJLTensorFType): + return False + return self.jltype == other.jltype + + def __hash__(self): + return hash(self.jltype) # TODO: Do we need the scipy and raw julia stuff class FinchJLTensor(_Display, EagerTensor): @@ -105,7 +114,7 @@ def _from_numpy( return jl.swizzle(jl.Tensor(lvl._obj), *order) @classmethod - def preprocess_order(cls, order: str, ndim: int) -> tuple[int, ...]: + def preprocess_order(cls, order: OrderType, ndim: int) -> tuple[int, ...]: if order == 'F': permutation = tuple(range(1, ndim + 1)) elif order == 'C': diff --git a/src/finch/typing.py b/src/finch/typing.py index fa305c9..7d64e03 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -13,3 +13,5 @@ DType = jc.AnyValue # represents jl.DataType Device = Literal["cpu"] | None + +OrderType = Literal["C", "F"] | tuple[int, ...] | None \ No newline at end of file diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 6c51fa5..d5438e2 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -25,14 +25,15 @@ import operator from finchlite import ftype from finchlite.algebra import overwrite, promote_min -from finchlite.compile import ExtentFType, dimension, BufferizedNDArray +from finchlite.compile import ExtentFType, dimension from finchlite.codegen import NumpyBuffer from finch.compiler import FinchJLCompiler +from finch.tensor import FinchJLTensor # Dummy data to obtain the bufferized ND array type a = np.zeros(dtype=np.float64, shape=(3, 3)) -a_format = ftype(BufferizedNDArray.from_numpy(a)) +a_format = ftype(FinchJLTensor(a)) @pytest.mark.parametrize( From 605d7e2c5ae1afed3b2a6ffd69a2f27e729e1434 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 11 Feb 2026 14:36:12 -0500 Subject: [PATCH 08/81] feat: improved call method in finch kernel --- .gitignore | 1 + src/finch/compiler.py | 14 +++-------- src/finch/tensor.py | 26 +++++++++---------- tests/test_compiler.py | 57 ++++++++++++++++++++++++++++++++++++------ 4 files changed, 67 insertions(+), 31 deletions(-) diff --git a/.gitignore b/.gitignore index 4896382..9167885 100644 --- a/.gitignore +++ b/.gitignore @@ -159,6 +159,7 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. .idea/ +.vscode/ # mac os .DS_Store diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 08f8d2e..2e85b1d 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -8,7 +8,7 @@ import operator -from juliacall import Main as jl +from .julia import jc, jl ops_map = {operator.add: "+", operator.mul: "*"} @@ -20,15 +20,9 @@ def __init__(self, func_name, jl_code): self.func_name = func_name jl.seval(self.jl_code) - # TODO: Switch back to (self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...] - def __call__(self, *args: tuple[FinchJLTensor, ...]): - argList = [] - for arg in args: - argList.append(f"arg{len(argList)}") - setattr(jl, argList[-1], arg) - - jl.seval(f"{self.func_name}({','.join(argList)})") - + def __call__(self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...]: + finch_fn = getattr(jl, self.func_name) + return tuple(finch_fn(*[arg._obj for arg in args])) class FinchJLLibrary(AssemblyLibrary): def __init__(self, kernel_dict): diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 335818f..944ceb0 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -15,6 +15,7 @@ sparse_formats_names, ) + class FinchJLTensorFType(TensorFType): def __init__(self, jltype, shape_type): # Julia type associated with the tensor @@ -23,13 +24,13 @@ def __init__(self, jltype, shape_type): def ndims(self) -> np.intp: return np.intp(jl.ndims(self.jltype)) - + def fill_value(self) -> Any: return jl.fill_value(self.jltype) - + def element_type(self) -> Any: return jl.eltype(self.jltype) - + def shape_type(self) -> tuple[type, ...]: return self._shape_type @@ -38,7 +39,7 @@ def __call__(self, shape: tuple) -> Tensor: def from_numpy(self, arr: np.ndarray) -> Tensor: return FinchJLTensor(arr) - + def __eq__(self, other): if not isinstance(other, FinchJLTensorFType): return False @@ -47,6 +48,7 @@ def __eq__(self, other): def __hash__(self): return hash(self.jltype) + # TODO: Do we need the scipy and raw julia stuff class FinchJLTensor(_Display, EagerTensor): def __init__( @@ -75,7 +77,7 @@ def __init__( "Either scalar, numpy, scipy.sparse or a raw julia object should " f"be provided. Found: {type(obj)}" ) - + @property def ftype(self): """ @@ -85,12 +87,12 @@ def ftype(self): for idx in self._shape: shape_type.append(type(idx)) return FinchJLTensorFType(jltype=jl.typeof(self._obj), shape_type=shape_type) - + @property def shape(self) -> tuple: """Shape of the tensor.""" return self._shape - + # TODO: do we need to have all the order stuff still? @classmethod def _from_numpy( @@ -112,15 +114,13 @@ def _from_numpy( for i in inv_order: lvl = Dense(lvl, arr.shape[i]) return jl.swizzle(jl.Tensor(lvl._obj), *order) - + @classmethod def preprocess_order(cls, order: OrderType, ndim: int) -> tuple[int, ...]: - if order == 'F': + if order == "F": permutation = tuple(range(1, ndim + 1)) - elif order == 'C': + elif order == "C": permutation = tuple(range(1, ndim + 1)[::-1]) else: - raise ValueError( - f"order must be 'C' or 'F'." - ) + raise ValueError(f"order must be 'C' or 'F'.") return permutation diff --git a/tests/test_compiler.py b/tests/test_compiler.py index d5438e2..f26585c 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -28,7 +28,7 @@ from finchlite.compile import ExtentFType, dimension from finchlite.codegen import NumpyBuffer -from finch.compiler import FinchJLCompiler +from finch.compiler import FinchJLCompiler, FinchJLKernel from finch.tensor import FinchJLTensor # Dummy data to obtain the bufferized ND array type @@ -36,6 +36,7 @@ a_format = ftype(FinchJLTensor(a)) +@pytest.mark.skip @pytest.mark.parametrize( "finch_ntn, julia_code", [ @@ -157,15 +158,13 @@ ), ), ), - Freeze(Slot("C_", a_format), Literal(operator.add)), - Repack( - Slot("C_", a_format), Variable("C", a_format) - ), - Return(Variable("C", a_format)), - ), + Freeze(Slot("C_", a_format), Literal(operator.add)), + Repack(Slot("C_", a_format), Variable("C", a_format)), + Return(Variable("C", a_format)), ), ), - ) + ), + ) ), """function matmul(C,A,B) @finch C .= 0.0 @@ -187,3 +186,45 @@ def test_finchjl_compiler(finch_ntn: Module, julia_code): compiler = FinchJLCompiler() library = compiler(finch_ntn) assert getattr(library, finch_ntn.children[0].name.name).jl_code == julia_code + + +@pytest.mark.parametrize( + "func_name, julia_prgm, args, expected_result", + [ + ( + "matmul", + """function matmul(C,A,B) + @finch C .= 0 + @finch begin + for i = _ + for k = _ + for j = _ + C[i,j] += *(A[i,k],B[k,j]) + end + end + end + end + return C +end""", + ( + FinchJLTensor(np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])), + FinchJLTensor(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])), + FinchJLTensor(np.array([[10, 11, 12], [13, 14, 15], [16, 17, 18]])), + ), + ( + FinchJLTensor( + np.array([[84, 90, 96], [201, 216, 231], [318, 342, 366]]) + ), + ), + ) + ], +) +def test_finchjl_kernel( + func_name: str, + julia_prgm: str, + args: tuple[FinchJLTensor, ...], + expected_result: tuple[FinchJLTensor, ...], +): + kernel = FinchJLKernel(func_name, julia_prgm) + result = kernel(*args) + assert result == expected_result From 0238adf10682ec2b3823d2c9fb91cd948e48d9e5 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 11 Feb 2026 14:43:21 -0500 Subject: [PATCH 09/81] feat: uncommenting code --- src/finch/_array_api_info.py | 174 +++++++++++++++++------------------ src/finch/dtypes.py | 122 ++++++++++++------------ src/finch/errors.py | 4 +- src/finch/tensor.py | 1 + 4 files changed, 151 insertions(+), 150 deletions(-) diff --git a/src/finch/_array_api_info.py b/src/finch/_array_api_info.py index b8d57c0..c4e3f5a 100644 --- a/src/finch/_array_api_info.py +++ b/src/finch/_array_api_info.py @@ -1,94 +1,94 @@ -# from . import dtypes -# from .typing import DType +from . import dtypes +from .typing import DType -# class __array_namespace_info__: -# def capabilities(self) -> dict[str, bool]: -# return { -# "boolean indexing": True, -# "data-dependent shapes": True, -# } +class __array_namespace_info__: + def capabilities(self) -> dict[str, bool]: + return { + "boolean indexing": True, + "data-dependent shapes": True, + } -# def default_device(self) -> str: -# return "cpu" + def default_device(self) -> str: + return "cpu" -# def default_dtypes(self, *, device: str | None = None) -> dict[str, DType]: -# if device not in ["cpu", None]: -# raise ValueError( -# f'Device not understood. Only "cpu" is allowed, but received: {device}' -# ) -# return { -# "real floating": dtypes.float64, -# "complex floating": dtypes.complex128, -# "integral": dtypes.int_, -# "indexing": dtypes.int_, -# } + def default_dtypes(self, *, device: str | None = None) -> dict[str, DType]: + if device not in ["cpu", None]: + raise ValueError( + f'Device not understood. Only "cpu" is allowed, but received: {device}' + ) + return { + "real floating": dtypes.float64, + "complex floating": dtypes.complex128, + "integral": dtypes.int_, + "indexing": dtypes.int_, + } -# _bool_dtypes = {"bool": dtypes.bool} -# _signed_integer_dtypes = { -# "int8": dtypes.int8, -# "int16": dtypes.int16, -# "int32": dtypes.int32, -# "int64": dtypes.int64, -# } -# _unsigned_integer_dtypes = { -# "uint8": dtypes.uint8, -# "uint16": dtypes.uint16, -# "uint32": dtypes.uint32, -# "uint64": dtypes.uint64, -# } -# _real_floating_dtypes = { -# "float32": dtypes.float32, -# "float64": dtypes.float64, -# } -# _complex_floating_dtypes = { -# "complex64": dtypes.complex64, -# "complex128": dtypes.complex128, -# } + _bool_dtypes = {"bool": dtypes.bool} + _signed_integer_dtypes = { + "int8": dtypes.int8, + "int16": dtypes.int16, + "int32": dtypes.int32, + "int64": dtypes.int64, + } + _unsigned_integer_dtypes = { + "uint8": dtypes.uint8, + "uint16": dtypes.uint16, + "uint32": dtypes.uint32, + "uint64": dtypes.uint64, + } + _real_floating_dtypes = { + "float32": dtypes.float32, + "float64": dtypes.float64, + } + _complex_floating_dtypes = { + "complex64": dtypes.complex64, + "complex128": dtypes.complex128, + } -# def dtypes( -# self, -# *, -# device: str | None = None, -# kind: str | tuple[str, ...] | None = None, -# ) -> dict[str, DType]: -# if device not in ["cpu", None]: -# raise ValueError( -# f'Device not understood. Only "cpu" is allowed, but received: {device}' -# ) -# if kind is None: -# return ( -# self._bool_dtypes -# | self._signed_integer_dtypes -# | self._unsigned_integer_dtypes -# | self._real_floating_dtypes -# | self._complex_floating_dtypes -# ) -# if kind == "bool": -# return self._bool_dtypes -# if kind == "signed integer": -# return self._signed_integer_dtypes -# if kind == "unsigned integer": -# return self._unsigned_integer_dtypes -# if kind == "integral": -# return self._signed_integer_dtypes | self._unsigned_integer_dtypes -# if kind == "real floating": -# return self._real_floating_dtypes -# if kind == "complex floating": -# return self._complex_floating_dtypes -# if kind == "numeric": -# return ( -# self._signed_integer_dtypes -# | self._unsigned_integer_dtypes -# | self._real_floating_dtypes -# | self._complex_floating_dtypes -# ) -# if isinstance(kind, tuple): -# res = {} -# for k in kind: -# res.update(self.dtypes(kind=k)) -# return res -# raise ValueError(f"unsupported kind: {kind!r}") + def dtypes( + self, + *, + device: str | None = None, + kind: str | tuple[str, ...] | None = None, + ) -> dict[str, DType]: + if device not in ["cpu", None]: + raise ValueError( + f'Device not understood. Only "cpu" is allowed, but received: {device}' + ) + if kind is None: + return ( + self._bool_dtypes + | self._signed_integer_dtypes + | self._unsigned_integer_dtypes + | self._real_floating_dtypes + | self._complex_floating_dtypes + ) + if kind == "bool": + return self._bool_dtypes + if kind == "signed integer": + return self._signed_integer_dtypes + if kind == "unsigned integer": + return self._unsigned_integer_dtypes + if kind == "integral": + return self._signed_integer_dtypes | self._unsigned_integer_dtypes + if kind == "real floating": + return self._real_floating_dtypes + if kind == "complex floating": + return self._complex_floating_dtypes + if kind == "numeric": + return ( + self._signed_integer_dtypes + | self._unsigned_integer_dtypes + | self._real_floating_dtypes + | self._complex_floating_dtypes + ) + if isinstance(kind, tuple): + res = {} + for k in kind: + res.update(self.dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") -# def devices(self) -> list[str]: -# return ["cpu"] + def devices(self) -> list[str]: + return ["cpu"] diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py index 3b6b4cc..2b95252 100644 --- a/src/finch/dtypes.py +++ b/src/finch/dtypes.py @@ -1,61 +1,61 @@ -# import builtins - -# import numpy as np - -# from .julia import jl - -# int_: jl.DataType = jl.Int -# int8: jl.DataType = jl.Int8 -# int16: jl.DataType = jl.Int16 -# int32: jl.DataType = jl.Int32 -# int64: jl.DataType = jl.Int64 -# uint: jl.DataType = jl.UInt -# uint8: jl.DataType = jl.UInt8 -# uint16: jl.DataType = jl.UInt16 -# uint32: jl.DataType = jl.UInt32 -# uint64: jl.DataType = jl.UInt64 -# float16: jl.DataType = jl.Float16 -# float32: jl.DataType = jl.Float32 -# float64: jl.DataType = jl.Float64 -# complex64: jl.DataType = jl.ComplexF32 -# complex128: jl.DataType = jl.ComplexF64 -# bool: jl.DataType = jl.Bool - -# number: jl.DataType = jl.Number -# complex: jl.DataType = jl.Complex -# integer: jl.DataType = jl.Integer -# abstract_float: jl.DataType = jl.AbstractFloat - -# jl_to_np_dtype = { -# int_: np.int_, -# int8: np.int8, -# int16: np.int16, -# int32: np.int32, -# int64: np.int64, -# uint: np.uint, -# uint8: np.uint8, -# uint16: np.uint16, -# uint32: np.uint32, -# uint64: np.uint64, -# float16: np.float16, -# float32: np.float32, -# float64: np.float64, -# complex64: np.complex64, -# complex128: np.complex128, -# bool: builtins.bool, -# None: None, -# } - - -# def finfo(dtype): -# return np.finfo(jl_to_np_dtype[dtype]) - - -# def iinfo(dtype): -# return np.iinfo(jl_to_np_dtype[dtype]) - - -# def can_cast(from_, to, /) -> builtins.bool: -# if hasattr(from_, "dtype"): -# from_ = from_.dtype -# return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) +import builtins + +import numpy as np + +from .julia import jl + +int_: jl.DataType = jl.Int +int8: jl.DataType = jl.Int8 +int16: jl.DataType = jl.Int16 +int32: jl.DataType = jl.Int32 +int64: jl.DataType = jl.Int64 +uint: jl.DataType = jl.UInt +uint8: jl.DataType = jl.UInt8 +uint16: jl.DataType = jl.UInt16 +uint32: jl.DataType = jl.UInt32 +uint64: jl.DataType = jl.UInt64 +float16: jl.DataType = jl.Float16 +float32: jl.DataType = jl.Float32 +float64: jl.DataType = jl.Float64 +complex64: jl.DataType = jl.ComplexF32 +complex128: jl.DataType = jl.ComplexF64 +bool: jl.DataType = jl.Bool + +number: jl.DataType = jl.Number +complex: jl.DataType = jl.Complex +integer: jl.DataType = jl.Integer +abstract_float: jl.DataType = jl.AbstractFloat + +jl_to_np_dtype = { + int_: np.int_, + int8: np.int8, + int16: np.int16, + int32: np.int32, + int64: np.int64, + uint: np.uint, + uint8: np.uint8, + uint16: np.uint16, + uint32: np.uint32, + uint64: np.uint64, + float16: np.float16, + float32: np.float32, + float64: np.float64, + complex64: np.complex64, + complex128: np.complex128, + bool: builtins.bool, + None: None, +} + + +def finfo(dtype): + return np.finfo(jl_to_np_dtype[dtype]) + + +def iinfo(dtype): + return np.iinfo(jl_to_np_dtype[dtype]) + + +def can_cast(from_, to, /) -> builtins.bool: + if hasattr(from_, "dtype"): + from_ = from_.dtype + return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) diff --git a/src/finch/errors.py b/src/finch/errors.py index 3a79630..c034e04 100644 --- a/src/finch/errors.py +++ b/src/finch/errors.py @@ -1,2 +1,2 @@ -# class PerformanceWarning(Warning): -# pass +class PerformanceWarning(Warning): + pass diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 944ceb0..1c5803c 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,6 +1,7 @@ import numpy as np from typing import Any from finchlite import EagerTensor, TensorFType, Tensor +from . import dtypes as jl_dtypes from .typing import OrderType, JuliaObj from .julia import jc, jl From 618823f63c8307a78a72b1ef20d4fc84416a8595 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 11 Feb 2026 15:57:01 -0500 Subject: [PATCH 10/81] feat: comments from meet --- src/finch/compiler.py | 2 ++ src/finch/tensor.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 2e85b1d..cde8248 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -13,6 +13,8 @@ ops_map = {operator.add: "+", operator.mul: "*"} +# https://github.com/finch-tensor/finch-tensor-lite/blob/main/tests/test_notation_interpreter.py + class FinchJLKernel(AssemblyKernel): def __init__(self, func_name, jl_code): # We store this code so that we can verify it in pytest diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 1c5803c..0ae01c0 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -16,6 +16,10 @@ sparse_formats_names, ) +# Singleton classes for levels types +# finch tensor lite, formatter stage +# level ftype without the need to create tthe object +# https://github.com/finch-tensor/finch-tensor-lite/blob/main/src/finchlite/autoschedule/formatter.py class FinchJLTensorFType(TensorFType): def __init__(self, jltype, shape_type): From 02323a445ee52e52a58448d81a75dbc1f3322ad3 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Thu, 19 Feb 2026 23:06:48 -0500 Subject: [PATCH 11/81] chore: readme correction --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 84b82cb..8273a2c 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ Finch uses [poetry](https://python-poetry.org/) for packaging. To install for development, clone the repository and run: ```bash -poetry install --with test +poetry install --extras test ``` to install the current project and dev dependencies. From 927f7ecfab7c9bdcb9413762f81289c2cc5742db Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 20 Feb 2026 01:45:20 -0500 Subject: [PATCH 12/81] chore: precommit --- src/finch/__init__.py | 362 ---------- src/finch/_array_api_info.py | 94 --- src/finch/dtypes.py | 61 -- src/finch/errors.py | 2 - src/finch/io.py | 15 - src/finch/levels.py | 135 ++-- src/finch/tensor-old.py | 577 --------------- src/finch/tensor.py | 126 +--- src/finch/typing.py | 8 +- tests/conftest.py | 35 - tests/data/matrix_1.ttx | 17 - tests/test_einsum.py | 1139 ------------------------------ tests/test_indexing.py | 125 ---- tests/test_io.py | 23 - tests/test_linalg.py | 36 - tests/test_ops.py | 440 ------------ tests/test_scipy_constructors.py | 123 ---- tests/test_sparse.py | 424 ----------- 18 files changed, 73 insertions(+), 3669 deletions(-) delete mode 100644 src/finch/_array_api_info.py delete mode 100644 src/finch/dtypes.py delete mode 100644 src/finch/errors.py delete mode 100644 src/finch/io.py delete mode 100644 src/finch/tensor-old.py delete mode 100644 tests/conftest.py delete mode 100644 tests/data/matrix_1.ttx delete mode 100644 tests/test_einsum.py delete mode 100644 tests/test_indexing.py delete mode 100644 tests/test_io.py delete mode 100644 tests/test_linalg.py delete mode 100644 tests/test_ops.py delete mode 100644 tests/test_scipy_constructors.py delete mode 100644 tests/test_sparse.py diff --git a/src/finch/__init__.py b/src/finch/__init__.py index eff4ae6..e69de29 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -1,362 +0,0 @@ -# from operator import ( -# abs as abs, -# ) -# from operator import ( -# add as add, -# ) -# from operator import ( -# and_ as bitwise_and, -# ) -# from operator import ( -# eq as equal, -# ) -# from operator import ( -# floordiv as floor_divide, -# ) -# from operator import ( -# ge as greater_equal, -# ) -# from operator import ( -# gt as greater, -# ) -# from operator import ( -# invert as bitwise_invert, -# ) -# from operator import ( -# le as less_equal, -# ) -# from operator import ( -# lshift as bitwise_left_shift, -# ) -# from operator import ( -# lt as less, -# ) -# from operator import ( -# matmul as matmul, -# ) -# from operator import ( -# mod as remainder, -# ) -# from operator import ( -# mul as multiply, -# ) -# from operator import ( -# ne as not_equal, -# ) -# from operator import ( -# neg as negative, -# ) -# from operator import ( -# or_ as bitwise_or, -# ) -# from operator import ( -# pos as positive, -# ) -# from operator import ( -# pow as pow, -# ) -# from operator import ( -# rshift as bitwise_right_shift, -# ) -# from operator import ( -# sub as subtract, -# ) -# from operator import ( -# truediv as divide, -# ) -# from operator import ( -# xor as bitwise_xor, -# ) - -# from numpy import ( -# e as e, -# ) -# from numpy import ( -# inf as inf, -# ) -# from numpy import ( -# nan as nan, -# ) -# from numpy import ( -# newaxis as newaxis, -# ) -# from numpy import ( -# pi as pi, -# ) - -# from . import linalg -# from ._array_api_info import __array_namespace_info__ -# from .compiled import ( -# DefaultScheduler, -# GalleyScheduler, -# compiled, -# compute, -# lazy, -# set_optimizer, -# ) -# from .dtypes import ( -# bool, -# can_cast, -# complex64, -# complex128, -# finfo, -# float16, -# float32, -# float64, -# iinfo, -# int8, -# int16, -# int32, -# int64, -# int_, -# uint, -# uint8, -# uint16, -# uint32, -# uint64, -# ) -# from .io import ( -# read, -# write, -# ) -# from .levels import ( -# Dense, -# DenseStorage, -# Element, -# Pattern, -# RepeatRLE, -# SparseByteMap, -# SparseCOO, -# SparseHash, -# SparseList, -# SparseVBL, -# Storage, -# ) -# from .tensor import ( -# SparseArray, -# FinchJLTensor, -# acos, -# acosh, -# all, -# any, -# arange, -# argmax, -# argmin, -# asarray, -# asin, -# asinh, -# astype, -# atan, -# atan2, -# atanh, -# ceil, -# conj, -# cos, -# cosh, -# diagonal, -# einop, -# einsum, -# empty, -# empty_like, -# exp, -# expand_dims, -# expm1, -# eye, -# floor, -# full, -# full_like, -# imag, -# isfinite, -# isinf, -# isnan, -# linspace, -# log, -# log1p, -# log2, -# log10, -# logaddexp, -# logical_and, -# logical_or, -# logical_xor, -# max, -# mean, -# min, -# moveaxis, -# nonzero, -# ones, -# ones_like, -# permute_dims, -# power, -# prod, -# random, -# real, -# reshape, -# round, -# sign, -# sin, -# sinh, -# sqrt, -# square, -# squeeze, -# std, -# sum, -# tan, -# tanh, -# tensordot, -# trunc, -# var, -# where, -# zeros, -# zeros_like, -# ) - -# __all__ = [ -# "DefaultScheduler", -# "Dense", -# "DenseStorage", -# "Element", -# "GalleyScheduler", -# "Pattern", -# "RepeatRLE", -# "SparseArray", -# "SparseByteMap", -# "SparseCOO", -# "SparseHash", -# "SparseList", -# "SparseVBL", -# "Storage", -# "FinchJLTensor", -# "__array_namespace_info__", -# "abs", -# "acos", -# "acosh", -# "add", -# "all", -# "any", -# "arange", -# "argmax", -# "argmin", -# "asarray", -# "asin", -# "asinh", -# "astype", -# "atan", -# "atan2", -# "atanh", -# "bitwise_and", -# "bitwise_invert", -# "bitwise_left_shift", -# "bitwise_or", -# "bitwise_right_shift", -# "bitwise_xor", -# "bool", -# "can_cast", -# "ceil", -# "compiled", -# "complex64", -# "complex128", -# "compute", -# "conj", -# "cos", -# "cosh", -# "diagonal", -# "divide", -# "e", -# "einop", -# "einsum", -# "empty", -# "empty_like", -# "equal", -# "exp", -# "expand_dims", -# "expm1", -# "eye", -# "finfo", -# "float16", -# "float32", -# "float64", -# "floor", -# "floor_divide", -# "full", -# "full_like", -# "greater", -# "greater_equal", -# "iinfo", -# "imag", -# "inf", -# "int8", -# "int16", -# "int32", -# "int64", -# "int_", -# "isfinite", -# "isinf", -# "isnan", -# "lazy", -# "less", -# "less_equal", -# "linalg", -# "linspace", -# "log", -# "log1p", -# "log2", -# "log10", -# "logaddexp", -# "logical_and", -# "logical_or", -# "logical_xor", -# "matmul", -# "max", -# "mean", -# "min", -# "moveaxis", -# "multiply", -# "nan", -# "negative", -# "newaxis", -# "nonzero", -# "not_equal", -# "ones", -# "ones_like", -# "permute_dims", -# "pi", -# "positive", -# "pow", -# "power", -# "prod", -# "random", -# "read", -# "real", -# "remainder", -# "reshape", -# "round", -# "set_optimizer", -# "sign", -# "sin", -# "sinh", -# "sqrt", -# "square", -# "squeeze", -# "std", -# "subtract", -# "sum", -# "tan", -# "tanh", -# "tensordot", -# "trunc", -# "uint", -# "uint8", -# "uint16", -# "uint32", -# "uint64", -# "var", -# "where", -# "write", -# "zeros", -# "zeros_like", -# ] - -# __array_api_version__: str = "2024.12" - - -from .compiler import FinchJLCompiler diff --git a/src/finch/_array_api_info.py b/src/finch/_array_api_info.py deleted file mode 100644 index c4e3f5a..0000000 --- a/src/finch/_array_api_info.py +++ /dev/null @@ -1,94 +0,0 @@ -from . import dtypes -from .typing import DType - - -class __array_namespace_info__: - def capabilities(self) -> dict[str, bool]: - return { - "boolean indexing": True, - "data-dependent shapes": True, - } - - def default_device(self) -> str: - return "cpu" - - def default_dtypes(self, *, device: str | None = None) -> dict[str, DType]: - if device not in ["cpu", None]: - raise ValueError( - f'Device not understood. Only "cpu" is allowed, but received: {device}' - ) - return { - "real floating": dtypes.float64, - "complex floating": dtypes.complex128, - "integral": dtypes.int_, - "indexing": dtypes.int_, - } - - _bool_dtypes = {"bool": dtypes.bool} - _signed_integer_dtypes = { - "int8": dtypes.int8, - "int16": dtypes.int16, - "int32": dtypes.int32, - "int64": dtypes.int64, - } - _unsigned_integer_dtypes = { - "uint8": dtypes.uint8, - "uint16": dtypes.uint16, - "uint32": dtypes.uint32, - "uint64": dtypes.uint64, - } - _real_floating_dtypes = { - "float32": dtypes.float32, - "float64": dtypes.float64, - } - _complex_floating_dtypes = { - "complex64": dtypes.complex64, - "complex128": dtypes.complex128, - } - - def dtypes( - self, - *, - device: str | None = None, - kind: str | tuple[str, ...] | None = None, - ) -> dict[str, DType]: - if device not in ["cpu", None]: - raise ValueError( - f'Device not understood. Only "cpu" is allowed, but received: {device}' - ) - if kind is None: - return ( - self._bool_dtypes - | self._signed_integer_dtypes - | self._unsigned_integer_dtypes - | self._real_floating_dtypes - | self._complex_floating_dtypes - ) - if kind == "bool": - return self._bool_dtypes - if kind == "signed integer": - return self._signed_integer_dtypes - if kind == "unsigned integer": - return self._unsigned_integer_dtypes - if kind == "integral": - return self._signed_integer_dtypes | self._unsigned_integer_dtypes - if kind == "real floating": - return self._real_floating_dtypes - if kind == "complex floating": - return self._complex_floating_dtypes - if kind == "numeric": - return ( - self._signed_integer_dtypes - | self._unsigned_integer_dtypes - | self._real_floating_dtypes - | self._complex_floating_dtypes - ) - if isinstance(kind, tuple): - res = {} - for k in kind: - res.update(self.dtypes(kind=k)) - return res - raise ValueError(f"unsupported kind: {kind!r}") - - def devices(self) -> list[str]: - return ["cpu"] diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py deleted file mode 100644 index 2b95252..0000000 --- a/src/finch/dtypes.py +++ /dev/null @@ -1,61 +0,0 @@ -import builtins - -import numpy as np - -from .julia import jl - -int_: jl.DataType = jl.Int -int8: jl.DataType = jl.Int8 -int16: jl.DataType = jl.Int16 -int32: jl.DataType = jl.Int32 -int64: jl.DataType = jl.Int64 -uint: jl.DataType = jl.UInt -uint8: jl.DataType = jl.UInt8 -uint16: jl.DataType = jl.UInt16 -uint32: jl.DataType = jl.UInt32 -uint64: jl.DataType = jl.UInt64 -float16: jl.DataType = jl.Float16 -float32: jl.DataType = jl.Float32 -float64: jl.DataType = jl.Float64 -complex64: jl.DataType = jl.ComplexF32 -complex128: jl.DataType = jl.ComplexF64 -bool: jl.DataType = jl.Bool - -number: jl.DataType = jl.Number -complex: jl.DataType = jl.Complex -integer: jl.DataType = jl.Integer -abstract_float: jl.DataType = jl.AbstractFloat - -jl_to_np_dtype = { - int_: np.int_, - int8: np.int8, - int16: np.int16, - int32: np.int32, - int64: np.int64, - uint: np.uint, - uint8: np.uint8, - uint16: np.uint16, - uint32: np.uint32, - uint64: np.uint64, - float16: np.float16, - float32: np.float32, - float64: np.float64, - complex64: np.complex64, - complex128: np.complex128, - bool: builtins.bool, - None: None, -} - - -def finfo(dtype): - return np.finfo(jl_to_np_dtype[dtype]) - - -def iinfo(dtype): - return np.iinfo(jl_to_np_dtype[dtype]) - - -def can_cast(from_, to, /) -> builtins.bool: - if hasattr(from_, "dtype"): - from_ = from_.dtype - return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) diff --git a/src/finch/errors.py b/src/finch/errors.py deleted file mode 100644 index c034e04..0000000 --- a/src/finch/errors.py +++ /dev/null @@ -1,2 +0,0 @@ -class PerformanceWarning(Warning): - pass diff --git a/src/finch/io.py b/src/finch/io.py deleted file mode 100644 index 4e10a72..0000000 --- a/src/finch/io.py +++ /dev/null @@ -1,15 +0,0 @@ -# from pathlib import Path - -# from .julia import jl -# from .tensor import Tensor - - -# def read(filename: Path | str) -> Tensor: -# fn = str(filename) -# julia_obj = jl.fread(fn) -# return Tensor(julia_obj) - - -# def write(filename: Path | str, tns: Tensor) -> None: -# fn = str(filename) -# jl.fwrite(fn, tns._obj) diff --git a/src/finch/levels.py b/src/finch/levels.py index 07d6142..284f6ae 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -1,110 +1,85 @@ -from .julia import jl -from .typing import DType, JuliaObj, OrderType - - -class _Display: - _obj: JuliaObj - - def __repr__(self): - return jl.sprint(jl.show, self._obj) - - def __str__(self): - return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) - +from abc import abstractmethod +from typing import Any -# LEVEL +import numpy as np +from finchlite import Tensor, TensorFType -class AbstractLevel(_Display): - pass - - -# core levels - - -class Dense(AbstractLevel): - def __init__(self, lvl, shape=None): - args = [lvl._obj] - if shape is not None: - args.append(shape) - self._obj = jl.Dense(*args) - +from .julia import jl +from .tensor import FinchJLTensor +from .typing import JuliaObj, number -class Element(AbstractLevel): - def __init__(self, fill_value, data=None): - args = [fill_value] - if data is not None: - args.append(data) - self._obj = jl.Element(*args) +class LevelFType(TensorFType): + def from_numpy(self, _) -> Tensor: + raise NotImplementedError -class Pattern(AbstractLevel): - def __init__(self): - self._obj = jl.Pattern() + def shape_type(self) -> tuple[type, ...]: + return tuple(self.element_type for _ in range(self.ndim)) -# advanced levels +class Element(LevelFType): + def __init__(self, fill_value: number): + self._fill_value = fill_value + def ndims(self) -> np.intp: + return np.intp(0) -class SparseList(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseList(lvl._obj) + def fill_value(self) -> Any: + return self._fill_value + def element_type(self) -> Any: + return type(self._fill_value) -class SparseByteMap(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseByteMap(lvl._obj) + def __call__(self, _) -> Tensor: + raise Exception("Cannot create an object of element type!") + def __eq__(self, other): + return isinstance(other, Element) and self._fill_value == other.fill_value -class RepeatRLE(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.RepeatRLE(lvl._obj) + def __hash__(self): + return hash((self.__class__.__name__, self._fill_value)) + def create_jl_obj(self) -> JuliaObj: + return jl.Element(self._fill_value) -class SparseVBL(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseVBL(lvl._obj) +class NestedLevelFType(LevelFType): + def __init__(self, lvl: LevelFType): + self.lvl = lvl -class SparseCOO(AbstractLevel): - def __init__(self, ndim, lvl): - self._obj = jl.SparseCOO[ndim](lvl._obj) + def ndims(self) -> np.intp: + return self.lvl.ndims + np.intp(1) + def fill_value(self) -> Any: + return self.lvl.fill_value -class SparseHash(AbstractLevel): - def __init__(self, ndim, lvl): - self._obj = jl.SparseHash[ndim](lvl._obj) + def element_type(self) -> Any: + return self.lvl.element_type + def __call__(self, shape: tuple) -> FinchJLTensor: + return FinchJLTensor(jl.Finch.Tensor(self.create_jl_obj(), shape)) -sparse_formats_names = ( - "SparseList", - "Sparse", - "SparseHash", - "SparseCOO", - "SparseRLE", - "SparseVBL", - "SparseBand", - "SparsePoint", - "SparseInterval", -) + def __eq__(self, other): + return type(other) is type(self) and self.lvl == other.lvl + def __hash__(self): + return hash((self.__class__.__name__, self.lvl.__hash__)) -# STORAGE + @abstractmethod + def create_jl_obj(self) -> JuliaObj: ... -class Storage: - def __init__(self, levels_descr: AbstractLevel, order: OrderType = None): - self.levels_descr = levels_descr - self.order = order if order is not None else "C" +class Dense(NestedLevelFType): + def create_jl_obj(self) -> JuliaObj: + return jl.Dense(self.lvl.create_jl_obj()) - def __str__(self) -> str: - return f"Storage(lvl={str(self.levels_descr)}, order={self.order})" +class SparseList(NestedLevelFType): + def create_jl_obj(self) -> JuliaObj: + return jl.SparseList(self.lvl.create_jl_obj()) -class DenseStorage(Storage): - def __init__(self, ndim: int, dtype: DType, order: OrderType = None): - lvl = Element(dtype(0)) - for _ in range(ndim): - lvl = Dense(lvl) - super().__init__(levels_descr=lvl, order=order) +class SparseByteMap(NestedLevelFType): + def create_jl_obj(self) -> JuliaObj: + return jl.SparseByteMap(self.lvl.create_jl_obj()) diff --git a/src/finch/tensor-old.py b/src/finch/tensor-old.py deleted file mode 100644 index afdffb1..0000000 --- a/src/finch/tensor-old.py +++ /dev/null @@ -1,577 +0,0 @@ -# from __future__ import annotations - -# import builtins -# import warnings -# from collections.abc import Callable, Iterable -# from typing import Any, Literal - -# import numpy as np -# from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple - -# from . import dtypes as jl_dtypes -# from .errors import PerformanceWarning -# from .julia import jc, jl -# from .levels import ( -# Dense, -# DenseStorage, -# Element, -# SparseCOO, -# SparseList, -# Storage, -# _Display, -# sparse_formats_names, -# ) -# from .typing import Device, DType, JuliaObj, OrderType, TupleOf3Arrays, spmatrix -# from finchlite import Tensor, TensorFType, EagerTensor - - -# class FinchJLTensorFType(TensorFType): -# def __init__(self, jltype): -# # Julia type associated with the tensor -# self.jltype = jltype - -# def ndims(self) -> np.intp: -# return np.intp(jl.ndims(self.jltype)) - -# def fill_value(self) -> Any: -# return jl.fill_value(self.jltype) - -# def element_type(self) -> Any: -# return jl.eltype(self.jltype) - -# # TODO: implement later -# def shape_type(self) -> tuple[type, ...]: -# ... - -# def __call__(self, shape: tuple) -> Tensor: -# ... - -# def from_numpy(self, arr: np.ndarray) -> Tensor: -# ... - -# class FinchJLTensor(_Display, EagerTensor): -# """ -# A wrapper class for Finch.Tensor and Finch.SwizzleArray. - -# Constructors -# ------------ -# FinchJLTensor(scipy.sparse.spmatrix) -# Construct a Tensor out of a `scipy.sparse` object. Supported formats are: `COO`, -# `CSC`, and `CSR`. -# FinchJLTensor(numpy.ndarray) -# Construct a Tensor out of a NumPy array object. This is a no-copy operation. -# FinchJLTensor(Storage) -# Initialize a Tensor with a `storage` description. `storage` can already hold -# data. -# FinchJLTensor(julia_object) -# Tensor created from a compatible raw Julia object. Must be a `Tensor`. -# This is a no-copy operation. - -# Parameters -# ---------- -# obj : np.ndarray or scipy.sparse or Storage or Finch.Tensor -# Input to construct a Tensor. It's a no-copy operation of for NumPy and -# SciPy input. For Storage it's levels' description with order. The order -# numbers the dimensions from the fastest to slowest. The leaf nodes have -# mode `0` and the root node has mode `n-1`. If the tensor was square of -# size `N`, then `N .^ order == strides`. Available options are "C" -# (row-major), "F" (column-major), or a custom order. Default: row-major. -# fill_value : np.number, optional -# Only used when `numpy.ndarray` or `scipy.sparse` is passed. -# copy : bool, optional -# If ``True``, then the object is copied. If ``None`` then the object is -# copied only if needed. For ``False`` it raises a ``ValueError`` if a -# copy cannot be avoided. Default: ``None``. - -# Returns -# ------- -# FinchJLTensor -# Python wrapper for Finch.jl `Tensor`. - -# Examples -# -------- -# >>> import numpy as np -# >>> import finch -# >>> arr2d = np.arange(6).reshape((2, 3)) -# >>> t1 = finch.FinchJLTensor(arr2d) -# >>> t1.todense() -# array([[0, 1, 2], -# [3, 4, 5]]) -# >>> np.shares_memory(t1.todense(), arr2d) -# True -# >>> storage = finch.Storage( -# ... finch.Dense(finch.SparseList(finch.Element(1))), order="C" -# ... ) -# >>> t2 = t1.to_storage(storage) -# >>> t2.todense() -# array([[0, 1, 2], -# [3, 4, 5]]) -# """ - -# def __init__( -# self, -# obj: np.ndarray | spmatrix | Storage | JuliaObj, -# /, -# *, -# fill_value: np.number | None = None, -# copy: bool | None = None, -# ): -# if isinstance(obj, int | float | complex | bool | list): -# if copy is False: -# raise ValueError( -# "copy=False isn't supported for scalar inputs and Python lists" -# ) -# obj = np.asarray(obj) -# if fill_value is None: -# fill_value = 0.0 - -# if _is_scipy_sparse_obj(obj): # scipy constructor -# jl_data = self._from_scipy_sparse(obj, fill_value=fill_value, copy=copy) -# self._obj = jl_data -# elif isinstance(obj, np.ndarray): # numpy constructor -# jl_data = self._from_numpy(obj, fill_value=fill_value, copy=copy) -# self._obj = jl_data -# elif isinstance(obj, Storage): # from-storage constructor -# if copy: -# self._raise_julia_copy_not_supported() -# order = self.preprocess_order( -# obj.order, self.get_lvl_ndim(obj.levels_descr._obj) -# ) -# self._obj = jl.swizzle(jl.Tensor(obj.levels_descr._obj), *order) -# elif jl.isa(obj, jl.Finch.Tensor): # raw-Julia-object constructors -# if copy: -# self._raise_julia_copy_not_supported() -# self._obj = jl.swizzle(obj, *tuple(range(1, jl.ndims(obj) + 1))) -# elif jl.isa(obj, jl.Finch.Tensor): -# if copy: -# self._raise_julia_copy_not_supported() -# self._obj = obj -# elif isinstance(obj, FinchJLTensor): -# self._obj = obj._obj -# else: -# raise ValueError( -# "Either scalar, numpy, scipy.sparse or a raw julia object should " -# f"be provided. Found: {type(obj)}" -# ) - -# @property -# def element_type(self): -# return jl.eltype(self._obj.body) - -# @property -# def dtype(self) -> DType: -# return jl.eltype(self._obj.body) - -# @property -# def ndim(self) -> int: -# return jl.ndims(self._obj) - -# @property -# def shape(self) -> tuple[int, ...]: -# return jl.size(self._obj) - -# @property -# def size(self) -> int: -# return np.prod(self.shape) - -# @property -# def fill_value(self) -> np.number: -# return jl.fill_value(self._obj) - -# @property -# def _is_dense(self) -> bool: -# lvl = self._obj.body.lvl -# for _ in self.shape: -# if not jl.isa(lvl, jl.Finch.Dense): -# return False -# lvl = lvl.lvl -# return True - -# @property -# def _order(self) -> tuple[int, ...]: -# return jl.typeof(self._obj).parameters[1] - -# @property -# def mT(self) -> Tensor: -# axes = list(range(self.ndim)) -# axes[-2], axes[-1] = axes[-1], axes[-2] -# axes = tuple(axes) -# return self.permute_dims(axes) - -# @property -# def device(self) -> str: -# return "cpu" - -# def to_device( -# self, device: Device, /, *, stream: int | Any | None = None -# ) -> Tensor: -# if device != "cpu": -# raise ValueError("Only `device='cpu'` is supported.") - -# return self - -# @classmethod -# def get_lvl_ndim(cls, lvl: JuliaObj) -> int: -# ndim = 0 -# while True: -# ndim += 1 -# lvl = lvl.lvl -# if jl.isa(lvl, jl.Finch.Element): -# break -# return ndim - -# def todense(self) -> np.ndarray: -# obj = self._obj - -# if self._is_dense: -# # don't materialize a dense finch tensor -# shape = jl.size(obj.body) -# dense_tensor = obj.body.lvl -# else: -# # create materialized dense array -# shape = jl.size(obj) -# dense_lvls = jl.Element(jc.convert(self.dtype, jl.fill_value(obj))) -# for _ in range(self.ndim): -# dense_lvls = jl.Dense(dense_lvls) -# dense_tensor = jl.Tensor(dense_lvls, obj).lvl # materialize - -# for _ in range(self.ndim): -# dense_tensor = dense_tensor.lvl - -# result = np.asarray(jl.reshape(dense_tensor.val, shape)) -# return result.transpose(self.get_order()) if self._is_dense else result - -# #TODO: Do we need? -# def permute_dims(self, axes: tuple[int, ...]) -> Tensor: -# axes = tuple(i + 1 for i in axes) -# new_obj = jl.permutedims(self._obj, axes) -# return Tensor(new_obj) - -# def to_storage(self, storage: Storage) -> Tensor: -# return Tensor(self._from_other_tensor(self, storage=storage)) - -# @classmethod -# def _from_other_tensor(cls, tensor: Tensor, storage: Storage) -> JuliaObj: -# order = cls.preprocess_order(storage.order, tensor.ndim) -# result = jl.copyto_b( -# jl.swizzle(jl.Tensor(storage.levels_descr._obj), *order), tensor._obj -# ) -# return jl.dropfills(result) if tensor._is_dense else result - -# @classmethod -# def _from_numpy( -# cls, arr: np.ndarray, fill_value: np.number, copy: bool | None = None -# ) -> JuliaObj: -# if copy: -# arr = arr.copy() -# order_char = "F" if np.isfortran(arr) else "C" -# order = cls.preprocess_order(order_char, arr.ndim) -# inv_order = tuple(i - 1 for i in jl.invperm(order)) - -# dtype = arr.dtype.type -# if ( -# dtype == np.bool_ -# ): # Fails with: Finch currently only supports isbits defaults -# dtype = jl_dtypes.bool -# fill_value = dtype(fill_value) -# lvl = Element(fill_value, arr.reshape(-1, order=order_char)) -# for i in inv_order: -# lvl = Dense(lvl, arr.shape[i]) -# return jl.swizzle(jl.Tensor(lvl._obj), *order) - -# @classmethod -# def from_scipy_sparse( -# cls, -# x, -# fill_value: np.number | None = None, -# copy: bool | None = None, -# ) -> Tensor: -# if not _is_scipy_sparse_obj(x): -# raise ValueError("{x} is not a SciPy sparse object.") -# return Tensor(x, fill_value=fill_value, copy=copy) - -# @classmethod -# def _from_scipy_sparse( -# cls, -# x, -# *, -# fill_value: np.number | None = None, -# copy: bool | None = None, -# ) -> JuliaObj: -# if copy is False and not ( -# x.format in ("coo", "csr", "csc") and x.has_canonical_format -# ): -# raise ValueError( -# "Unable to avoid copy while creating an array as requested." -# ) -# if x.format not in ("coo", "csr", "csc"): -# x = x.asformat("coo") -# if copy: -# x = x.copy() -# if not x.has_canonical_format: -# x.sum_duplicates() -# assert x.has_canonical_format - -# if x.format == "coo": -# return cls.construct_coo_jl_object( -# coords=(x.col, x.row), -# data=x.data, -# shape=x.shape[::-1], -# order=Tensor.row_major, -# fill_value=fill_value, -# ) -# if x.format == "csc": -# return cls.construct_csc_jl_object( -# arg=(x.data, x.indices, x.indptr), -# shape=x.shape, -# fill_value=fill_value, -# ) -# if x.format == "csr": -# return cls.construct_csr_jl_object( -# arg=(x.data, x.indices, x.indptr), -# shape=x.shape, -# fill_value=fill_value, -# ) -# raise ValueError(f"Unsupported SciPy format: {type(x)}") - -# @classmethod -# def construct_coo_jl_object( -# cls, coords, data, shape, order, fill_value=0.0 -# ) -> JuliaObj: -# assert len(coords) == 2 -# ndim = len(shape) -# order = cls.preprocess_order(order, ndim) - -# lvl = jl.Element(data.dtype.type(fill_value), data) -# ptr = jl.Vector[jl.Int]([1, len(data) + 1]) -# tbl = tuple(jl.PlusOneVector(arr) for arr in coords) - -# return jl.swizzle(jl.Tensor(jl.SparseCOO[ndim](lvl, shape, ptr, tbl)), *order) - -# @classmethod -# def construct_coo( -# cls, coords, data, shape, order=row_major, fill_value=0.0 -# ) -> Tensor: -# return Tensor( -# cls.construct_coo_jl_object(coords, data, shape, order, fill_value) -# ) - -# @staticmethod -# def _construct_compressed2d_jl_object( -# arg: TupleOf3Arrays, -# shape: tuple[int, ...], -# order: tuple[int, ...], -# fill_value: np.number = 0.0, -# ) -> JuliaObj: -# assert isinstance(arg, tuple) and len(arg) == 3 -# assert len(shape) == 2 - -# data, indices, indptr = arg -# dtype = data.dtype.type -# indices = jl.PlusOneVector(indices) -# indptr = jl.PlusOneVector(indptr) - -# lvl = jl.Element(dtype(fill_value), data) -# return jl.swizzle( -# jl.Tensor( -# jl.Dense(jl.SparseList(lvl, shape[0], indptr, indices), shape[1]) -# ), -# *order, -# ) - -# @classmethod -# def construct_csc_jl_object( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> JuliaObj: -# return cls._construct_compressed2d_jl_object( -# arg=arg, shape=shape, order=(1, 2), fill_value=fill_value -# ) - -# @classmethod -# def construct_csc( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> Tensor: -# return Tensor(cls.construct_csc_jl_object(arg, shape, fill_value)) - -# @classmethod -# def construct_csr_jl_object( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> JuliaObj: -# return cls._construct_compressed2d_jl_object( -# arg=arg, shape=shape[::-1], order=(2, 1), fill_value=fill_value -# ) - -# @classmethod -# def construct_csr( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> Tensor: -# return Tensor(cls.construct_csr_jl_object(arg, shape, fill_value)) - -# @staticmethod -# def construct_csf_jl_object( -# arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> JuliaObj: -# assert isinstance(arg, tuple) and len(arg) == 3 - -# data, indices_list, indptr_list = arg -# dtype = data.dtype.type - -# assert len(indices_list) == len(shape) - 1 -# assert len(indptr_list) == len(shape) - 1 - -# indices_list = [jl.PlusOneVector(i) for i in indices_list] -# indptr_list = [jl.PlusOneVector(i) for i in indptr_list] - -# lvl = jl.Element(dtype(fill_value), data) -# for size, indices, indptr in zip( -# shape[:-1], indices_list, indptr_list, strict=False -# ): -# lvl = jl.SparseList(lvl, size, indptr, indices) - -# return jl.swizzle( -# jl.Tensor(jl.Dense(lvl, shape[-1])), *range(1, len(shape) + 1) -# ) - -# @classmethod -# def construct_csf( -# cls, arg: TupleOf3Arrays, shape: tuple[int, ...], fill_value: np.number = 0.0 -# ) -> Tensor: -# return Tensor(cls.construct_csf_jl_object(arg, shape, fill_value)) - -# def to_scipy_sparse(self, accept_fv=None): -# import scipy.sparse as sp - -# if accept_fv is None: -# accept_fv = [0] -# elif not isinstance(accept_fv, Iterable): -# accept_fv = [accept_fv] - -# if self.ndim != 2: -# raise ValueError( -# "Can only convert a 2-dimensional array to a Scipy sparse matrix." -# ) -# if not builtins.any(_eq_scalars(self.fill_value, fv) for fv in accept_fv): -# raise ValueError( -# f"Can only convert arrays with {accept_fv} fill-values " -# "to a Scipy sparse matrix." -# ) -# order = self.get_order() -# body = self._obj.body - -# if str(jl.typeof(body.lvl).name.name) == "SparseCOOLevel": -# data = np.asarray(body.lvl.lvl.val) -# coords = body.lvl.tbl -# row, col = coords[::-1] if order == (1, 0) else coords -# row, col = np.asarray(row) - 1, np.asarray(col) - 1 -# return sp.coo_matrix((data, (row, col)), shape=self.shape) - -# if ( -# str(jl.typeof(body.lvl).name.name) == "DenseLevel" -# and str(jl.typeof(body.lvl.lvl).name.name) == "SparseListLevel" -# ): -# data = np.asarray(body.lvl.lvl.lvl.val) -# indices = np.asarray(body.lvl.lvl.idx) - 1 -# indptr = np.asarray(body.lvl.lvl.ptr) - 1 -# sp_class = sp.csr_matrix if order == (1, 0) else sp.csc_matrix -# return sp_class((data, indices, indptr), shape=self.shape) -# if ( -# jl.typeof(body.lvl).name.name in sparse_formats_names -# or jl.typeof(body.lvl.lvl).name.name in sparse_formats_names -# ): -# storage = Storage(SparseCOO(self.ndim, Element(self.fill_value)), order) -# return self.to_storage(storage).to_scipy_sparse() -# raise ValueError("Tensor can't be converted to scipy.sparse object.") - -# @staticmethod -# def _raise_julia_copy_not_supported() -> None: -# raise ValueError("copy=True isn't supported for Julia object inputs") - -# def __array_namespace__(self, *, api_version: str | None = None) -> Any: -# if api_version is None: -# api_version = "2024.12" - -# if api_version not in {"2021.12", "2022.12", "2023.12", "2024.12"}: -# raise ValueError(f'"{api_version}" Array API version not supported.') -# import finch - -# return finch - - -# def random(shape, density=0.01, random_state=None): -# args = [*shape, density] -# if random_state is not None: -# if isinstance(random_state, np.random.Generator): -# seed = random_state.integers(np.iinfo(np.int32).max) -# else: -# seed = random_state -# rng = jl.Random.default_rng() -# jl.Random.seed_b(rng, seed) -# args = [rng] + args -# return Tensor(jl.fsprand(*args)) - - -# def asarray( -# obj, -# /, -# *, -# dtype: DType | None = None, -# format: str | None = None, -# fill_value: np.number | None = None, -# device: Device | None = None, -# copy: bool | None = None, -# ) -> Tensor: -# if format not in {"coo", "csr", "csc", "csf", "dense", None}: -# raise ValueError(f"{format} format not supported.") -# _validate_device(device) -# tensor = ( -# obj -# if isinstance(obj, Tensor) -# else Tensor(obj, fill_value=fill_value, copy=copy) -# ) -# if format is not None: -# if copy is False: -# raise ValueError( -# "Unable to avoid copy while creating an array as requested." -# ) -# order = tensor.get_order() -# if format == "coo": -# storage = Storage(SparseCOO(tensor.ndim, Element(tensor.fill_value)), order) -# elif format == "csr": -# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (2, 1)) -# elif format == "csc": -# storage = Storage(Dense(SparseList(Element(tensor.fill_value))), (1, 2)) -# elif format == "csf": -# storage = Element(tensor.fill_value) -# for _ in range(tensor.ndim - 1): -# storage = SparseList(storage) -# storage = Storage(Dense(storage), order) -# elif format == "dense": -# storage = DenseStorage(tensor.ndim, tensor.dtype, order) -# tensor = tensor.to_storage(storage) - -# if dtype is not None: -# return astype(tensor, dtype, copy=copy) -# return tensor - - -# def reshape( -# x: Tensor, /, shape: tuple[int, ...], *, copy: bool | None = None -# ) -> Tensor: -# if copy is False: -# raise ValueError("Unable to avoid copy during reshape.") -# # TODO: https://github.com/finch-tensor/Finch.jl/issues/743 -# # Revert to `jl.reshape` implementation once aforementioned -# # issue is solved. -# warnings.warn( -# "`reshape` densified the input tensor.", PerformanceWarning, stacklevel=2 -# ) -# arr = x.todense() -# arr = arr.reshape(shape) -# return Tensor(arr) - - -# def _validate_device(device: Device) -> None: -# if device not in {"cpu", None}: -# raise ValueError( -# f'Device not understood. Only "cpu" is allowed, but received: {device}' -# ) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 0ae01c0..d700f9a 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,131 +1,37 @@ -import numpy as np -from typing import Any -from finchlite import EagerTensor, TensorFType, Tensor -from . import dtypes as jl_dtypes +from finchlite import EagerTensor -from .typing import OrderType, JuliaObj -from .julia import jc, jl -from .levels import ( - Dense, - DenseStorage, - Element, - SparseCOO, - SparseList, - Storage, - _Display, - sparse_formats_names, -) +from .julia import jl +from .typing import JuliaObj # Singleton classes for levels types # finch tensor lite, formatter stage # level ftype without the need to create tthe object # https://github.com/finch-tensor/finch-tensor-lite/blob/main/src/finchlite/autoschedule/formatter.py -class FinchJLTensorFType(TensorFType): - def __init__(self, jltype, shape_type): - # Julia type associated with the tensor - self.jltype = jltype - self._shape_type = shape_type - def ndims(self) -> np.intp: - return np.intp(jl.ndims(self.jltype)) +class _Display: + _obj: JuliaObj - def fill_value(self) -> Any: - return jl.fill_value(self.jltype) + def __repr__(self): + return jl.sprint(jl.show, self._obj) - def element_type(self) -> Any: - return jl.eltype(self.jltype) + def __str__(self): + return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) - def shape_type(self) -> tuple[type, ...]: - return self._shape_type - def __call__(self, shape: tuple) -> Tensor: - return FinchJLTensor(np.ones(shape=shape)) - - def from_numpy(self, arr: np.ndarray) -> Tensor: - return FinchJLTensor(arr) - - def __eq__(self, other): - if not isinstance(other, FinchJLTensorFType): - return False - return self.jltype == other.jltype - - def __hash__(self): - return hash(self.jltype) - - -# TODO: Do we need the scipy and raw julia stuff class FinchJLTensor(_Display, EagerTensor): - def __init__( - self, - obj: np.ndarray, - /, - *, - fill_value: np.number | None = None, - copy: bool | None = None, - ): - if isinstance(obj, int | float | complex | bool | list): - if copy is False: - raise ValueError( - "copy=False isn't supported for scalar inputs and Python lists" - ) - obj = np.asarray(obj) - if fill_value is None: - fill_value = 0.0 - - if isinstance(obj, np.ndarray): # numpy constructor - jl_data = self._from_numpy(obj, fill_value=fill_value, copy=copy) - self._shape = obj.shape - self._obj = jl_data + def __init__(self, obj: jl.Finch.Tensor): + if isinstance(obj, jl.Finch.Tensor): + self._obj = obj else: - raise ValueError( - "Either scalar, numpy, scipy.sparse or a raw julia object should " - f"be provided. Found: {type(obj)}" - ) + raise ValueError(f"Raw julia object expected. Found: {type(obj)}") + # TODO: figure out a way to walk through the levels and return the ftype @property def ftype(self): - """ - Returns the ftype of the buffer, which is a BufferizedNDArrayFType. - """ - shape_type = [] - for idx in self._shape: - shape_type.append(type(idx)) - return FinchJLTensorFType(jltype=jl.typeof(self._obj), shape_type=shape_type) + """Returns the ftype of the buffer""" @property def shape(self) -> tuple: """Shape of the tensor.""" - return self._shape - - # TODO: do we need to have all the order stuff still? - @classmethod - def _from_numpy( - cls, arr: np.ndarray, fill_value: np.number, copy: bool | None = None - ) -> JuliaObj: - if copy: - arr = arr.copy() - order_char = "F" if np.isfortran(arr) else "C" - order = cls.preprocess_order(order_char, arr.ndim) - inv_order = tuple(i - 1 for i in jl.invperm(order)) - - dtype = arr.dtype.type - if ( - dtype == np.bool_ - ): # Fails with: Finch currently only supports isbits defaults - dtype = jl_dtypes.bool - fill_value = dtype(fill_value) - lvl = Element(fill_value, arr.reshape(-1, order=order_char)) - for i in inv_order: - lvl = Dense(lvl, arr.shape[i]) - return jl.swizzle(jl.Tensor(lvl._obj), *order) - - @classmethod - def preprocess_order(cls, order: OrderType, ndim: int) -> tuple[int, ...]: - if order == "F": - permutation = tuple(range(1, ndim + 1)) - elif order == "C": - permutation = tuple(range(1, ndim + 1)[::-1]) - else: - raise ValueError(f"order must be 'C' or 'F'.") - return permutation + return self.obj.shape diff --git a/src/finch/typing.py b/src/finch/typing.py index 7d64e03..125086e 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -1,11 +1,7 @@ -from typing import Literal, Any - -import numpy as np +from typing import Any, Literal import juliacall as jc -TupleOf3Arrays = tuple[np.ndarray, np.ndarray, np.ndarray] - spmatrix = Any JuliaObj = jc.AnyValue @@ -14,4 +10,4 @@ Device = Literal["cpu"] | None -OrderType = Literal["C", "F"] | tuple[int, ...] | None \ No newline at end of file +number = int | float | bool | complex diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index 455e172..0000000 --- a/tests/conftest.py +++ /dev/null @@ -1,35 +0,0 @@ -import pytest - -import numpy as np - - -@pytest.fixture -def rng(): - return np.random.default_rng(42) - - -@pytest.fixture -def arr1d(): - return np.arange(100) - - -@pytest.fixture -def arr2d(): - return np.array( - [ - [0, 0, 3, 2, 0], - [1, 0, 0, 1, 0], - [0, 5, 0, 0, 0], - ] - ) - - -@pytest.fixture -def arr3d(): - return np.array( - [ - [[0, 1, 0, 0], [1, 0, 0, 3]], - [[4, 0, -1, 0], [2, 2, 0, 0]], - [[0, 0, 0, 0], [1, 5, 0, 3]], - ] - ) diff --git a/tests/data/matrix_1.ttx b/tests/data/matrix_1.ttx deleted file mode 100644 index 9c4de1c..0000000 --- a/tests/data/matrix_1.ttx +++ /dev/null @@ -1,17 +0,0 @@ -%%MatrixMarket matrix coordinate integer general -3 5 15 -1 1 0 -2 1 1 -3 1 0 -1 2 0 -2 2 0 -3 2 5 -1 3 3 -2 3 0 -3 3 0 -1 4 2 -2 4 1 -3 4 0 -1 5 0 -2 5 0 -3 5 0 diff --git a/tests/test_einsum.py b/tests/test_einsum.py deleted file mode 100644 index 02df717..0000000 --- a/tests/test_einsum.py +++ /dev/null @@ -1,1139 +0,0 @@ -import pytest - -import numpy as np - -import finch - - -@pytest.fixture -def rng(): - return np.random.default_rng(42) - - -def test_basic_addition_with_transpose(rng): - """Test basic addition with transpose""" - A = rng.random((5, 5)) - B = rng.random((5, 5)) - - C = finch.einop("C[i,j] = A[i,j] + B[j,i]", A=A, B=B).todense() - C_ref = A + B.T - - assert np.allclose(C, C_ref) - - -def test_matrix_multiplication(rng): - """Test matrix multiplication using += (increment/accumulation)""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - C = finch.einop("C[i,j] += A[i,k] * B[k,j]", A=A, B=B).todense() - C_ref = A @ B - - assert np.allclose(C, C_ref) - - -def test_element_wise_multiplication(rng): - """Test element-wise multiplication""" - A = rng.random((4, 4)) - B = rng.random((4, 4)) - - C = finch.einop("C[i,j] = A[i,j] * B[i,j]", A=A, B=B).todense() - C_ref = A * B - - assert np.allclose(C, C_ref) - - -def test_sum_reduction(rng): - """Test sum reduction using +=""" - A = rng.random((3, 4)) - - C = finch.einop("C[i] += A[i,j]", A=A).todense() - C_ref = np.sum(A, axis=1) - - assert np.allclose(C, C_ref) - - -def test_maximum_reduction(rng): - """Test maximum reduction using max=""" - A = rng.random((3, 4)) - - C = finch.einop("C[i] max= A[i,j]", A=A).todense() - C_ref = np.max(A, axis=1) - - assert np.allclose(C, C_ref) - - -def test_outer_product(rng): - """Test outer product""" - A = rng.random(3) - B = rng.random(4) - - C = finch.einop("C[i,j] = A[i] * B[j]", A=A, B=B).todense() - C_ref = np.outer(A, B) - - assert np.allclose(C, C_ref) - - -def test_batch_matrix_multiplication(rng): - """Test batch matrix multiplication using +=""" - A = rng.random((2, 3, 4)) - B = rng.random((2, 4, 5)) - - C = finch.einop("C[b,i,j] += A[b,i,k] * B[b,k,j]", A=A, B=B).todense() - C_ref = np.matmul(A, B) - - assert np.allclose(C, C_ref) - - -def test_minimum_reduction(rng): - """Test minimum reduction using min=""" - A = rng.random((3, 4)) - - C = finch.einop("C[i] min= A[i,j]", A=A).todense() - C_ref = np.min(A, axis=1) - - assert np.allclose(C, C_ref) - - -@pytest.mark.parametrize("axis", [(0, 2, 1), (3, 0, 1), (1, 0, 3, 2), (1, 0, 3, 2)]) -@pytest.mark.parametrize( - "idxs", - [ - ("i", "j", "k", "l"), - ("l", "j", "k", "i"), - ("l", "k", "j", "i"), - ], -) -def test_swizzle_in(rng, axis, idxs): - """Test transpositions with einop""" - A = rng.random((4, 4, 4, 4)) - - jdxs = [idxs[p] for p in axis] - xp_idxs = ", ".join(idxs) - np_idxs = "".join(idxs) - xp_jdxs = ", ".join(jdxs) - np_jdxs = "".join(jdxs) - - C = finch.einop(f"C[{xp_jdxs}] += A[{xp_idxs}]", A=A).todense() - C_ref = np.einsum(f"{np_idxs}->{np_jdxs}", A) - - assert np.allclose(C, C_ref) - - -def test_operator_precedence_arithmetic(rng): - """Test that arithmetic operator precedence follows Python rules""" - A = rng.random((3, 3)) - B = rng.random((3, 3)) - C = rng.random((3, 3)) - - # Test: A + B * C should be A + (B * C), not (A + B) * C - result = finch.einop("D[i,j] = A[i,j] + B[i,j] * C[i,j]", A=A, B=B, C=C).todense() - expected = A + (B * C) - - assert np.allclose(result, expected) - - -def test_operator_precedence_power_and_multiplication(rng): - """Test that power has higher precedence than multiplication""" - A = rng.random((3, 3)) + 1 # Add 1 to avoid numerical issues with powers - - # Test: A * A ** 2 should be A * (A ** 2), not (A * A) ** 2 - result = finch.einop("B[i,j] = A[i,j] * A[i,j] ** 2", A=A).todense() - expected = A * (A**2) - - assert np.allclose(result, expected) - - -def test_operator_precedence_addition_and_multiplication(rng): - """Test complex arithmetic precedence: A + B * C ** 2""" - A = rng.random((3, 3)) - B = rng.random((3, 3)) - C = rng.random((3, 3)) + 1 # Add 1 to avoid numerical issues - - # Test: A + B * C ** 2 should be A + (B * (C ** 2)) - result = finch.einop( - "D[i,j] = A[i,j] + B[i,j] * C[i,j] ** 2", A=A, B=B, C=C - ).todense() - expected = A + (B * (C**2)) - - assert np.allclose(result, expected) - - -def test_operator_precedence_logical_and_or(rng): - """Test that 'and' has higher precedence than 'or'""" - A = rng.random((3, 3)) > 0.3 # Boolean-like arrays - B = rng.random((3, 3)) > 0.3 - C = rng.random((3, 3)) > 0.3 - - # Test: A or B and C should be A or (B and C), not (A or B) and C - result = finch.einop( - "D[i,j] = A[i,j] or B[i,j] and C[i,j]", A=A, B=B, C=C - ).todense() - expected = np.logical_or(A, np.logical_and(B, C)) - - assert np.allclose(result, expected) - - -def test_operator_precedence_bitwise_operations(rng): - """Test bitwise operator precedence. - - | has lower precedence than ^ which has lower than & - """ - # Use integer arrays for bitwise operations - A = rng.integers(0, 8, size=(3, 3)) - B = rng.integers(0, 8, size=(3, 3)) - C = rng.integers(0, 8, size=(3, 3)) - D = rng.integers(0, 8, size=(3, 3)) - - # Test: A | B ^ C & D should be A | (B ^ (C & D)) - result = finch.einop( - "E[i,j] = A[i,j] | B[i,j] ^ C[i,j] & D[i,j]", A=A, B=B, C=C, D=D - ).todense() - expected = A | (B ^ (C & D)) - - assert np.allclose(result, expected) - - -def test_operator_precedence_shift_operations(rng): - """Test shift operator precedence with arithmetic""" - # Use small integer arrays to avoid overflow in shifts - A = rng.integers(1, 4, size=(3, 3)) - - # Test: A << 1 + 1 should be A << (1 + 1), not (A << 1) + 1 - # Since shift has lower precedence than addition - result = finch.einop("B[i,j] = A[i,j] << 1 + 1", A=A).todense() - expected = A << (1 + 1) # A << 2 - - assert np.allclose(result, expected) - - -def test_operator_precedence_comparison_with_arithmetic(rng): - """Test that arithmetic has higher precedence than comparison""" - A = rng.random((3, 3)) - B = rng.random((3, 3)) - C = rng.random((3, 3)) - - # Test: A + B == C should be (A + B) == C, not A + (B == C) - result = finch.einop("D[i,j] = A[i,j] + B[i,j] == C[i,j]", A=A, B=B, C=C).todense() - expected = (A + B) == C - - assert np.allclose(result, expected) - - -def test_operator_precedence_with_parentheses(rng): - """Test that parentheses override operator precedence""" - A = rng.random((3, 3)) - B = rng.random((3, 3)) - C = rng.random((3, 3)) - - # Test: (A + B) * C should be different from A + B * C - result_with_parens = finch.einop( - "D[i,j] = (A[i,j] + B[i,j]) * C[i,j]", A=A, B=B, C=C - ).todense() - result_without_parens = finch.einop( - "E[i,j] = A[i,j] + B[i,j] * C[i,j]", A=A, B=B, C=C - ).todense() - - expected_with_parens = (A + B) * C - expected_without_parens = A + (B * C) - - assert np.allclose(result_with_parens, expected_with_parens) - assert np.allclose(result_without_parens, expected_without_parens) - - # Verify they're different (unless by coincidence) - if not np.allclose(expected_with_parens, expected_without_parens): - assert not np.allclose(result_with_parens, result_without_parens) - - -def test_operator_precedence_unary_operators(rng): - """Test unary operator precedence""" - A = rng.random((3, 3)) - 0.5 # Some negative values - - # Test: -A ** 2 should be -(A ** 2), not (-A) ** 2 - result = finch.einop("B[i,j] = -A[i,j] ** 2", A=A).todense() - expected = -(A**2) - - assert np.allclose(result, expected) - - -def test_numeric_literals(rng): - """Test that numeric literals work correctly""" - A = rng.random((3, 3)) - - # Test simple addition with literal - result = finch.einop("B[i,j] = A[i,j] + 1", A=A).todense() - expected = A + 1 - - assert np.allclose(result, expected) - - # Test complex expression with literals - result2 = finch.einop("C[i,j] = A[i,j] * 2 + 3", A=A).todense() - expected2 = A * 2 + 3 - - assert np.allclose(result2, expected2) - - -def test_comparison_chaining(rng): - """Test that comparison chaining works like Python. - - a < b < c becomes (a < b) and (b < c) - """ - A = rng.random((3, 3)) * 10 # Scale to get variety in comparisons - B = rng.random((3, 3)) * 10 - C = rng.random((3, 3)) * 10 - - # Test: A < B < C should be (A < B) and (B < C), not (A < B) < C - result = finch.einop("D[i,j] = A[i,j] < B[i,j] < C[i,j]", A=A, B=B, C=C).todense() - expected = np.logical_and(A < B, B < C) - - assert np.allclose(result, expected) - - -def test_comparison_chaining_three_way(rng): - """Test three-way comparison chaining with different operators""" - A = np.array([[1, 2], [3, 4]]) - B = np.array([[2, 3], [4, 5]]) - C = np.array([[3, 4], [5, 6]]) - - # Test: A <= B < C should be (A <= B) and (B < C) - result = finch.einop("D[i,j] = A[i,j] <= B[i,j] < C[i,j]", A=A, B=B, C=C).todense() - expected = np.logical_and(A <= B, B < C) - - assert np.allclose(result, expected) - - -def test_comparison_chaining_four_way(rng): - """Test four-way comparison chaining""" - A = np.array([[1]]) - B = np.array([[2]]) - C = np.array([[3]]) - D = np.array([[4]]) - - # Test: A < B < C < D should be ((A < B) and (B < C)) and (C < D) - result = finch.einop( - "E[i,j] = A[i,j] < B[i,j] < C[i,j] < D[i,j]", A=A, B=B, C=C, D=D - ).todense() - expected = np.logical_and(np.logical_and(A < B, B < C), C < D) - - assert np.allclose(result, expected) - - -def test_single_comparison_vs_chained(rng): - """Test that single comparison and chained comparison work differently""" - A = np.array([[2]]) - B = np.array([[3]]) - C = np.array([[1]]) # Intentionally make C < A to show difference - - # Single comparison: A < B should be True - result_single = finch.einop("D[i,j] = A[i,j] < B[i,j]", A=A, B=B).todense() - expected_single = A < B - - # Chained comparison: A < B < C should be (A < B) and (B < C) - # = True and False = False - result_chained = finch.einop( - "E[i,j] = A[i,j] < B[i,j] < C[i,j]", A=A, B=B, C=C - ).todense() - expected_chained = np.logical_and(A < B, B < C) - - assert np.allclose(result_single, expected_single) - assert np.allclose(result_chained, expected_chained) - - # Verify they're different - assert not np.allclose(result_single, result_chained) - - -def test_alphanumeric_tensor_names(rng): - """Test that tensor names with numbers work correctly""" - A1 = rng.random((2, 2)) - B2 = rng.random((2, 2)) - C3_test = rng.random((2, 2)) - - # Test basic arithmetic with alphanumeric names - result = finch.einop( - "result_1[i,j] = A1[i,j] + B2[i,j] * C3_test[i,j]", - A1=A1, - B2=B2, - C3_test=C3_test, - ).todense() - expected = A1 + (B2 * C3_test) - - assert np.allclose(result, expected) - - # Test comparison chaining with alphanumeric names - X1 = np.array([[1, 2]]) - Y2 = np.array([[3, 4]]) - Z3 = np.array([[5, 6]]) - - result2 = finch.einop( - "chain_result[i,j] = X1[i,j] < Y2[i,j] < Z3[i,j]", X1=X1, Y2=Y2, Z3=Z3 - ).todense() - expected2 = np.logical_and(X1 < Y2, Y2 < Z3).astype(float) - - assert np.allclose(result2, expected2) - - -def test_bool_literals(rng): - """Test that boolean literals work correctly""" - A = rng.random((2, 2)) > 0.5 - - # Test True literal - result_true = finch.einop("B[i,j] = A[i,j] and True", A=A).todense() - expected_true = np.logical_and(A, True).astype(float) - assert np.allclose(result_true, expected_true) - - # Test False literal - result_false = finch.einop("C[i,j] = A[i,j] or False", A=A).todense() - expected_false = np.logical_or(A, False).astype(float) - assert np.allclose(result_false, expected_false) - - # Test boolean operations with literals - A_bool = rng.random((2, 2)) > 0.5 - result_and = finch.einop( - "D[i,j] = A_bool[i,j] and True and False", A_bool=A_bool - ).todense() - expected_and = np.logical_and(np.logical_and(A_bool, True), False) - assert np.allclose(result_and, expected_and) - - -def test_int_literals(rng): - """Test that integer literals work correctly""" - A = rng.random((2, 2)) - - # Test positive integer - result_pos = finch.einop("B[i,j] = A[i,j] + 42", A=A).todense() - expected_pos = A + 42 - assert np.allclose(result_pos, expected_pos) - - # Test negative integer - result_neg = finch.einop("C[i,j] = A[i,j] * -5", A=A).todense() - expected_neg = A * (-5) - assert np.allclose(result_neg, expected_neg) - - # Test zero - result_zero = finch.einop("D[i,j] = A[i,j] + 0", A=A).todense() - expected_zero = A + 0 - assert np.allclose(result_zero, expected_zero) - - # Test large integer - result_large = finch.einop("E[i,j] = A[i,j] + 123456789", A=A).todense() - expected_large = A + 123456789 - assert np.allclose(result_large, expected_large) - - -def test_float_literals(rng): - """Test that float literals work correctly""" - A = rng.random((2, 2)) - - # Test positive float - result_pos = finch.einop("B[i,j] = A[i,j] + 3.14159", A=A).todense() - expected_pos = A + 3.14159 - assert np.allclose(result_pos, expected_pos) - - # Test negative float - result_neg = finch.einop("C[i,j] = A[i,j] * -2.71828", A=A).todense() - expected_neg = A * (-2.71828) - assert np.allclose(result_neg, expected_neg) - - # Test scientific notation - result_sci = finch.einop("D[i,j] = A[i,j] + 1.5e-3", A=A).todense() - expected_sci = A + 1.5e-3 - assert np.allclose(result_sci, expected_sci) - - # Test very small float - result_small = finch.einop("E[i,j] = A[i,j] + 0.000001", A=A).todense() - expected_small = A + 0.000001 - assert np.allclose(result_small, expected_small) - - -def test_complex_literals(rng): - """Test that complex literals work correctly""" - A = rng.random((2, 2)).astype(complex) # Use complex arrays - - # Test complex with real and imaginary parts - result_complex = finch.einop("B[i,j] = A[i,j] + (3+4j)", A=A).todense() - expected_complex = A + (3 + 4j) - assert np.allclose(result_complex, expected_complex) - - # Test pure imaginary - result_imag = finch.einop("C[i,j] = A[i,j] * 2j", A=A).todense() - expected_imag = A * 2j - assert np.allclose(result_imag, expected_imag) - - # Test complex with negative parts - result_neg = finch.einop("D[i,j] = A[i,j] + (-1-2j)", A=A).todense() - expected_neg = A + (-1 - 2j) - assert np.allclose(result_neg, expected_neg) - - -def test_mixed_literal_types(rng): - """Test expressions mixing different literal types""" - A = rng.random((2, 2)) - - # Test int + float - result_int_float = finch.einop("B[i,j] = A[i,j] + 5 + 3.14", A=A).todense() - expected_int_float = A + 5 + 3.14 - assert np.allclose(result_int_float, expected_int_float) - - # Test operator precedence with literals - result_precedence = finch.einop("C[i,j] = A[i,j] + 2 * 3", A=A).todense() - expected_precedence = A + (2 * 3) # Should be A + 6, not (A + 2) * 3 - assert np.allclose(result_precedence, expected_precedence) - - # Test power with literals - result_power = finch.einop("D[i,j] = A[i,j] + 2 ** 3", A=A).todense() - expected_power = A + (2**3) # Should be A + 8 - assert np.allclose(result_power, expected_power) - - -def test_literal_edge_cases1(rng): - """Test edge cases with literals""" - A = rng.random((2, 2)) - - # Test multiple literals in sequence - result_multi = finch.einop("B[i,j] = A[i,j] + 1 + 2 + 3", A=A).todense() - expected_multi = A + 1 + 2 + 3 # Should be A + 6 - assert np.allclose(result_multi, expected_multi) - - -def test_literal_edge_cases2(rng): - """Test edge cases with literals""" - A = rng.random((2, 2)) - # Test literals in comparisons - result_comp = finch.einop("C[i,j] = A[i,j] > 0.5", A=A).todense() - expected_comp = (A > 0.5).astype(float) - assert np.allclose(result_comp, expected_comp) - - -def test_literal_edge_cases3(rng): - """Test edge cases with literals""" - A = rng.random((2, 2)) - # Test literals with parentheses - result_parens = finch.einop("D[i,j] = A[i,j] * (2 + 3)", A=A).todense() - expected_parens = A * (2 + 3) # Should be A * 5 - assert np.allclose(result_parens, expected_parens) - - -# ============================================================================= -# Einsum tests comparing finch.einsum to np.einsum -# ============================================================================= - - -class TestEinsumImplicitMode: - """Test einsum in implicit mode (no -> in subscripts)""" - - @pytest.mark.skip(reason="Repeated indices not yet supported") - def test_trace(self, rng): - """Test trace of a matrix""" - A = rng.random((5, 5)) - - result = finch.einsum("ii", A) - expected = np.einsum("ii", A) - - assert np.allclose(result, expected) - assert np.allclose(result, np.trace(A)) - - def test_element_wise_multiplication(self, rng): - """Test element-wise multiplication""" - A = rng.random((4, 3)) - B = rng.random((4, 3)) - - result = finch.einsum("ij,ij", A, B).todense() - expected = np.einsum("ij,ij", A, B) - - assert np.allclose(result, expected) - assert np.allclose(result, np.sum(A * B)) - - def test_matrix_multiplication(self, rng): - """Test matrix multiplication""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - result = finch.einsum("ij,jk", A, B).todense() - expected = np.einsum("ij,jk", A, B) - - assert np.allclose(result, expected) - assert np.allclose(result, A @ B) - - def test_transpose(self, rng): - """Test transpose via einsum""" - A = rng.random((3, 4)) - - result = finch.einsum("ji", A).todense() - expected = np.einsum("ji", A) - - assert np.allclose(result, expected) - assert np.allclose(result, A.T) - - def test_vector_inner_product(self, rng): - """Test vector inner product""" - a = rng.random(5) - b = rng.random(5) - - result = finch.einsum("i,i", a, b).todense() - expected = np.einsum("i,i", a, b) - - assert np.allclose(result, expected) - assert np.allclose(result, np.inner(a, b)) - - def test_vector_outer_product(self, rng): - """Test vector outer product""" - a = rng.random(3) - b = rng.random(4) - - result = finch.einsum("i,j", a, b).todense() - expected = np.einsum("i,j", a, b) - - assert np.allclose(result, expected) - assert np.allclose(result, np.outer(a, b)) - - def test_tensor_contraction(self, rng): - """Test tensor contraction""" - A = rng.random((3, 4, 5)) - B = rng.random((4, 3, 2)) - - result = finch.einsum("ijk,jil", A, B).todense() - expected = np.einsum("ijk,jil", A, B) - - assert np.allclose(result, expected) - - def test_bilinear_transformation(self, rng): - """Test bilinear transformation""" - A = rng.random((3, 4)) - B = rng.random((4, 5, 6)) - C = rng.random((6, 7)) - - result = finch.einsum("ij,jkl,lm", A, B, C).todense() - expected = np.einsum("ij,jkl,lm", A, B, C) - - assert np.allclose(result, expected) - - -class TestEinsumExplicitMode: - """Test einsum in explicit mode (with -> in subscripts)""" - - @pytest.mark.skip(reason="Repeated indices not yet supported") - def test_extract_diagonal(self, rng): - """Test extracting diagonal""" - A = rng.random((5, 5)) - - result = finch.einsum("ii->i", A) - expected = np.einsum("ii->i", A) - - assert np.allclose(result, expected) - assert np.allclose(result, np.diag(A)) - - def test_sum_over_axis(self, rng): - """Test sum over specific axis""" - A = rng.random((4, 5)) - - # Sum over axis 1 - result = finch.einsum("ij->i", A).todense() - expected = np.einsum("ij->i", A) - - assert np.allclose(result, expected) - assert np.allclose(result, np.sum(A, axis=1)) - - # Sum over axis 0 - result2 = finch.einsum("ij->j", A).todense() - expected2 = np.einsum("ij->j", A) - - assert np.allclose(result2, expected2) - assert np.allclose(result2, np.sum(A, axis=0)) - - def test_total_sum(self, rng): - """Test total sum""" - A = rng.random((3, 4)) - - result = finch.einsum("ij->", A).todense() - expected = np.einsum("ij->", A) - - assert np.allclose(result, expected) - assert np.allclose(result, np.sum(A)) - - def test_transpose_explicit(self, rng): - """Test transpose with explicit output""" - A = rng.random((3, 4)) - - result = finch.einsum("ij->ji", A).todense() - expected = np.einsum("ij->ji", A) - - assert np.allclose(result, expected) - assert np.allclose(result, A.T) - - def test_matrix_vector_explicit(self, rng): - """Test matrix-vector multiplication with explicit output""" - A = rng.random((3, 4)) - b = rng.random(4) - - result = finch.einsum("ij,j->i", A, b).todense() - expected = np.einsum("ij,j->i", A, b) - - assert np.allclose(result, expected) - assert np.allclose(result, A @ b) - - def test_custom_contraction(self, rng): - """Test custom tensor contraction with explicit output""" - A = rng.random((2, 3, 4)) - B = rng.random((4, 5)) - - result = finch.einsum("ijk,kl->ijl", A, B).todense() - expected = np.einsum("ijk,kl->ijl", A, B) - - assert np.allclose(result, expected) - - def test_reorder_axes(self, rng): - """Test reordering axes with explicit output""" - A = rng.random((2, 3, 4, 5)) - - result = finch.einsum("ijkl->ljik", A).todense() - expected = np.einsum("ijkl->ljik", A) - - assert np.allclose(result, expected) - - -class TestEinsumVariableOrdering: - """Test different variable orderings in einsum""" - - def test_alphabetical_ordering_implicit(self, rng): - """Test that implicit mode respects alphabetical ordering""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - # Standard order - result1 = finch.einsum("ij,jk", A, B).todense() - expected1 = np.einsum("ij,jk", A, B) - assert np.allclose(result1, expected1) - - # Different variable names but same meaning - result2 = finch.einsum("ab,bc", A, B).todense() - expected2 = np.einsum("ab,bc", A, B) - assert np.allclose(result2, expected2) - assert np.allclose(result1, result2) - - def test_non_alphabetical_ordering(self, rng): - """Test non-alphabetical variable ordering""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - # Non-alphabetical order should transpose result in implicit mode - result = finch.einsum("ij,jl", A, B).todense() - expected = np.einsum("ij,jl", A, B) - - assert np.allclose(result, expected) - # This should be different from alphabetical ij,jk due to l < k - - def test_variable_ordering_explicit_override(self, rng): - """Test that explicit mode overrides alphabetical ordering""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - # Force specific output order - result = finch.einsum("ij,jk->ki", A, B).todense() - expected = np.einsum("ij,jk->ki", A, B) - - assert np.allclose(result, expected) - # This should be transpose of standard matrix multiplication - assert np.allclose(result, (A @ B).T) - - @pytest.mark.skip(reason="Repeated indices not yet supported") - def test_repeated_indices(self, rng): - """Test repeated indices for diagonal operations""" - A = rng.random((4, 4, 4)) - - # Repeated index in same tensor - result = finch.einsum("iji", A) - expected = np.einsum("iji", A) - - assert np.allclose(result, expected) - - def test_mixed_variable_types(self, rng): - """Test mixing different variable names""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - # Mix letters from different parts of alphabet - result = finch.einsum("az,zb", A, B).todense() - expected = np.einsum("az,zb", A, B) - - assert np.allclose(result, expected) - - -class TestEinsumSpecialSyntax: - """Test einsum special syntax: einsum(op1, op1inds, op2, op2inds, ...)""" - - def test_alternative_syntax_matrix_mult(self, rng): - """Test alternative syntax for matrix multiplication""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - # Standard syntax - result1 = finch.einsum("ij,jk", A, B).todense() - - # Alternative syntax - result2 = finch.einsum(A, [0, 1], B, [1, 2]).todense() - - expected = np.einsum(A, [0, 1], B, [1, 2]) - - assert np.allclose(result1, result2) - assert np.allclose(result2, expected) - - def test_alternative_syntax_with_output(self, rng): - """Test alternative syntax with explicit output specification""" - A = rng.random((3, 4)) - B = rng.random((4, 5)) - - # Explicit output indices - result = finch.einsum(A, [0, 1], B, [1, 2], [0, 2]).todense() - expected = np.einsum(A, [0, 1], B, [1, 2], [0, 2]) - - assert np.allclose(result, expected) - assert np.allclose(result, A @ B) - - def test_alternative_syntax_transpose(self, rng): - """Test alternative syntax for transpose""" - A = rng.random((3, 4)) - - # Implicit transpose - result1 = finch.einsum(A, [1, 0]).todense() - expected1 = np.einsum(A, [1, 0]) - - assert np.allclose(result1, expected1) - assert np.allclose(result1, A.T) - - # Explicit transpose - result2 = finch.einsum(A, [0, 1], [1, 0]).todense() - expected2 = np.einsum(A, [0, 1], [1, 0]) - - assert np.allclose(result2, expected2) - assert np.allclose(result2, A.T) - - @pytest.mark.skip(reason="Repeated indices not yet supported") - def test_alternative_syntax_trace(self, rng): - """Test alternative syntax for trace""" - A = rng.random((5, 5)) - - # Trace using alternative syntax - result = finch.einsum(A, [0, 0]) - expected = np.einsum(A, [0, 0]) - - assert np.allclose(result, expected) - assert np.allclose(result, np.trace(A)) - - @pytest.mark.skip(reason="Repeated indices not yet supported") - def test_alternative_syntax_diagonal(self, rng): - """Test alternative syntax for diagonal extraction""" - A = rng.random((5, 5)) - - result = finch.einsum(A, [0, 0], [0]) - expected = np.einsum(A, [0, 0], [0]) - - assert np.allclose(result, expected) - assert np.allclose(result, np.diag(A)) - - def test_alternative_syntax_sum(self, rng): - """Test alternative syntax for sum operations""" - A = rng.random((3, 4)) - - # Sum over axis 1 - result1 = finch.einsum(A, [0, 1], [0]).todense() - expected1 = np.einsum(A, [0, 1], [0]) - - assert np.allclose(result1, expected1) - assert np.allclose(result1, np.sum(A, axis=1)) - - # Total sum - result2 = finch.einsum(A, [0, 1], []).todense() - expected2 = np.einsum(A, [0, 1], []) - - assert np.allclose(result2, expected2) - assert np.allclose(result2, np.sum(A)) - - def test_alternative_syntax_outer_product(self, rng): - """Test alternative syntax for outer product""" - a = rng.random(3) - b = rng.random(4) - - result = finch.einsum(a, [0], b, [1]).todense() - expected = np.einsum(a, [0], b, [1]) - - assert np.allclose(result, expected) - assert np.allclose(result, np.outer(a, b)) - - def test_alternative_syntax_complex_contraction(self, rng): - """Test alternative syntax for complex tensor contraction""" - A = rng.random((2, 3, 4)) - B = rng.random((4, 5, 6)) - C = rng.random((6, 2)) - - result = finch.einsum(A, [0, 1, 2], B, [2, 3, 4], C, [4, 0], [1, 3]).todense() - expected = np.einsum(A, [0, 1, 2], B, [2, 3, 4], C, [4, 0], [1, 3]) - - assert np.allclose(result, expected) - - -class TestEinsumEdgeCases: - """Test edge cases and special scenarios""" - - def test_single_tensor_operations(self, rng): - """Test operations on single tensors""" - A = rng.random((3, 4, 5)) - - # Identity operation - result1 = finch.einsum("ijk", A).todense() - expected1 = np.einsum("ijk", A) - assert np.allclose(result1, expected1) - assert np.allclose(result1, A) - - # Permute dimensions - result2 = finch.einsum("ikj", A).todense() - expected2 = np.einsum("ikj", A) - assert np.allclose(result2, expected2) - - def test_scalar_operations(self, rng): - """Test operations involving scalars""" - scalar = rng.random() - A = rng.random((3, 4)) - - # Scalar multiplication (broadcasting) - result = finch.einsum(",ij", scalar, A).todense() - expected = np.einsum(",ij", scalar, A) - - assert np.allclose(result, expected) - assert np.allclose(result, scalar * A) - - def test_empty_dimensions(self, rng): - """Test with empty dimensions""" - A = rng.random((0, 3)) - B = rng.random((3, 4)) - - result = finch.einsum("ij,jk", A, B).todense() - expected = np.einsum("ij,jk", A, B) - - assert np.allclose(result, expected) - assert result.shape == (0, 4) - - def test_1d_tensors(self, rng): - """Test operations on 1D tensors""" - a = rng.random(5) - b = rng.random(5) - - # Element-wise product and sum - result = finch.einsum("i,i", a, b).todense() - expected = np.einsum("i,i", a, b) - - assert np.allclose(result, expected) - assert np.isscalar(result) or result.shape == () - - def test_high_dimensional(self, rng): - """Test high-dimensional tensors""" - A = rng.random((2, 2, 2, 2, 2)) - B = rng.random((2, 2, 2, 2, 2)) - - result = finch.einsum("abcde,abcde", A, B).todense() - expected = np.einsum("abcde,abcde", A, B) - - assert np.allclose(result, expected) - - -class TestEinsumEllipses: - """Test einsum with ellipses (...) notation""" - - def test_basic_ellipses(self, rng): - """Test basic ellipses usage for identity operations""" - A = rng.random((3, 4, 5)) - - # Identity with ellipses - result = finch.einsum("...", A).todense() - expected = np.einsum("...", A) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_with_named_indices(self, rng): - """Test ellipses combined with named indices""" - A = rng.random((2, 3, 4, 5)) - - # Sum over last dimension, keeping others - result = finch.einsum("...i->...", A).todense() - expected = np.einsum("...i->...", A) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_transpose(self, rng): - """Test ellipses with transpose operations""" - A = rng.random((2, 3, 4, 5)) - - # Transpose last two dimensions - result = finch.einsum("...ij->...ji", A).todense() - expected = np.einsum("...ij->...ji", A) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_matrix_multiply(self, rng): - """Test batch matrix multiplication with ellipses""" - A = rng.random((2, 3, 4, 5)) - B = rng.random((2, 3, 5, 6)) - - # Batch matrix multiplication - result = finch.einsum("...ij,...jk->...ik", A, B).todense() - expected = np.einsum("...ij,...jk->...ik", A, B) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_different_batch_dims(self, rng): - """Test ellipses with different numbers of batch dimensions""" - A = rng.random((2, 3, 4)) # 1 batch dim + 2x4 matrix - B = rng.random((5, 2, 4, 6)) # 2 batch dims + 4x6 matrix - - # Broadcasting should work - result = finch.einsum("...ij,...jk->...ik", A, B).todense() - expected = np.einsum("...ij,...jk->...ik", A, B) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - @pytest.mark.skip(reason="Repeated indices not yet supported") - def test_ellipses_trace(self, rng): - """Test computing trace with ellipses""" - A = rng.random((2, 3, 4, 4)) - - # Trace of last two dimensions for each batch - result = finch.einsum("...ii->...", A) - expected = np.einsum("...ii->...", A) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - @pytest.mark.skip(reason="Repeated indices not yet supported") - def test_ellipses_diagonal(self, rng): - """Test extracting diagonal with ellipses""" - A = rng.random((2, 3, 4, 4)) - - # Extract diagonal of last two dimensions - result = finch.einsum("...ii->...i", A) - expected = np.einsum("...ii->...i", A) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_element_wise(self, rng): - """Test element-wise operations with ellipses""" - A = rng.random((2, 3, 4, 5)) - B = rng.random((2, 3, 4, 5)) - - # Element-wise multiplication - result = finch.einsum("...,...->...", A, B).todense() - expected = np.einsum("...,...->...", A, B) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_sum_product(self, rng): - """Test sum of element-wise product with ellipses""" - A = rng.random((2, 3, 4, 5)) - B = rng.random((2, 3, 4, 5)) - - # Sum of element-wise product - result = finch.einsum("...,...", A, B).todense() - expected = np.einsum("...,...", A, B) - - assert np.allclose(result, expected) - - def test_ellipses_outer_product(self, rng): - """Test outer product with ellipses""" - A = rng.random((2, 3)) - B = rng.random((2, 5)) - - # Outer product with batch dimensions - result = finch.einsum("...i,...j->...ij", A, B).todense() - expected = np.einsum("...i,...j->...ij", A, B) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_multiple_contractions(self, rng): - """Test multiple contractions with ellipses""" - A = rng.random((2, 3, 4, 5)) - B = rng.random((2, 3, 5, 6)) - C = rng.random((2, 3, 6, 7)) - - # Chain of matrix multiplications - result = finch.einsum("...ij,...jk,...kl->...il", A, B, C).todense() - expected = np.einsum("...ij,...jk,...kl->...il", A, B, C) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_broadcasting_edge_cases(self, rng): - """Test edge cases with broadcasting and ellipses""" - # Test with single dimension arrays - A = rng.random((1, 3, 4)) - B = rng.random((2, 1, 4, 5)) - - result = finch.einsum("...ij,...jk->...ik", A, B).todense() - expected = np.einsum("...ij,...jk->...ik", A, B) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_with_scalars(self, rng): - """Test ellipses operations with scalar inputs""" - A = rng.random((2, 3, 4)) - scalar = 2.5 - - # Multiply tensor by scalar using ellipses - result = finch.einsum("...,...->...", A, scalar).todense() - expected = np.einsum("...,...->...", A, scalar) - - assert np.allclose(result, expected) - assert result.shape == expected.shape - - def test_ellipses_reduction_patterns(self, rng): - """Test various reduction patterns with ellipses""" - A = rng.random((2, 3, 4, 5, 6)) - - # Sum over last dimension - result1 = finch.einsum("...i->...", A).todense() - expected1 = np.einsum("...i->...", A) - assert np.allclose(result1, expected1) - - # Sum over last two dimensions - result2 = finch.einsum("...ij->...", A).todense() - expected2 = np.einsum("...ij->...", A) - assert np.allclose(result2, expected2) - - # Sum over specific dimensions while keeping others - result3 = finch.einsum("...ijk->...ik", A).todense() - expected3 = np.einsum("...ijk->...ik", A) - assert np.allclose(result3, expected3) - - -@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.complex64, np.complex128]) -class TestEinsumDataTypes: - """Test einsum with different data types""" - - def test_matrix_multiplication_dtypes(self, rng, dtype): - """Test matrix multiplication with different dtypes""" - A = rng.random((3, 4)).astype(dtype) - B = rng.random((4, 5)).astype(dtype) - - result = finch.einsum("ij,jk", A, B).todense() - expected = np.einsum("ij,jk", A, B) - - assert np.allclose(result, expected) - # Check dtype preservation (may depend on implementation) - - def test_complex_operations(self, rng, dtype): - """Test operations with complex numbers""" - if not np.issubdtype(dtype, np.complexfloating): - pytest.skip("Test only for complex dtypes") - - A = (rng.random((3, 3)) + 1j * rng.random((3, 3))).astype(dtype) - - result = finch.einsum("ij", A).todense() - expected = np.einsum("ij", A) - - assert np.allclose(result, expected) diff --git a/tests/test_indexing.py b/tests/test_indexing.py deleted file mode 100644 index 865a2a7..0000000 --- a/tests/test_indexing.py +++ /dev/null @@ -1,125 +0,0 @@ -import pytest - -import numpy as np -from numpy.testing import assert_equal - -import juliacall as jc - -import finch - - -@pytest.mark.parametrize( - "index", - [ - ..., - 40, - (32,), - slice(None), - slice(30, 60, 3), - -10, - slice(None, -10, -2), - (None, slice(None)), - ], -) -@pytest.mark.parametrize("order", ["C", "F"]) -def test_indexing_1d(arr1d, index, order): - arr = np.array(arr1d, order=order) - arr_finch = finch.Tensor(arr) - - actual = arr_finch[index] - expected = arr[index] - - if isinstance(actual, finch.Tensor): - actual = actual.todense() - - assert_equal(actual, expected) - - -@pytest.mark.parametrize( - "index", - [ - ..., - 0, - (2,), - (2, 3), - slice(None), - (..., slice(0, 4, 2)), - (-1, slice(-1, None, -1)), - (None, slice(None), slice(None)), - ], -) -@pytest.mark.parametrize("order", ["C", "F"]) -def test_indexing_2d(arr2d, index, order): - arr = np.array(arr2d, order=order) - arr_finch = finch.Tensor(arr) - - actual = arr_finch[index] - expected = arr[index] - - if isinstance(actual, finch.Tensor): - actual = actual.todense() - - assert_equal(actual, expected) - - -@pytest.mark.parametrize( - "index", - [ - (0, 1, 2), - (1, 0, 0), - (0, 1), - 1, - 2, - (2, slice(None), 3), - (slice(None), 0), - slice(None), - (0, slice(None), slice(1, 4, 2)), - (0, 1, ...), - (..., 1), - (0, ..., 1), - ..., - (..., slice(1, 4, 2)), - (slice(None, None, -1), slice(None, None, -1), slice(None, None, -1)), - (slice(None, -1, 1), slice(-1, None, -1), slice(4, 1, -1)), - (-1, 0, 0), - (0, -1, -2), - ([1, 2], 0, slice(3, None, -1)), - (0, slice(1, 0, -1), 0), - (slice(None), None, slice(None), slice(None)), - (slice(None), slice(None), slice(None), None), - ], -) -@pytest.mark.parametrize( - "levels_descr", - [ - finch.Dense(finch.Dense(finch.Dense(finch.Element(0)))), - finch.Dense(finch.SparseList(finch.SparseList(finch.Element(0)))), - ], -) -@pytest.mark.parametrize("order", ["C", "F"]) -def test_indexing_3d(arr3d, index, levels_descr, order): - arr = np.array(arr3d, order=order) - storage = finch.Storage(levels_descr, order=order) - arr_finch = finch.Tensor(arr).to_storage(storage) - - actual = arr_finch[index] - expected = arr[index] - - if isinstance(actual, finch.Tensor): - actual = actual.todense() - - assert_equal(actual, expected) - - -def test_lazy_none_ellipsis(arr3d): - arr_finch = finch.lazy(finch.Tensor(arr3d)) - assert_equal(finch.compute(arr_finch[..., None]).todense(), arr3d[..., None]) - - with pytest.raises( - jc.JuliaError, - match=( - "Cannot index a lazy tensor with more or fewer `:` dims than it had " - "original dims." - ), - ): - arr_finch[None, :] diff --git a/tests/test_io.py b/tests/test_io.py deleted file mode 100644 index 4782f8d..0000000 --- a/tests/test_io.py +++ /dev/null @@ -1,23 +0,0 @@ -from numpy.testing import assert_equal - -import finch - -base_path = "tests/data" - - -def test_read(arr2d): - tns = finch.read(f"{base_path}/matrix_1.ttx") - - assert_equal(tns.todense(), arr2d) - - -def test_write(tmp_path, arr2d): - tns = finch.asarray(arr2d) - finch.write(tmp_path / "tmp.ttx", tns) - - with open(f"{base_path}/matrix_1.ttx") as f: - expected = f.read() - with open(tmp_path / "tmp.ttx") as f: - actual = f.read() - - assert actual == expected diff --git a/tests/test_linalg.py b/tests/test_linalg.py deleted file mode 100644 index 220776d..0000000 --- a/tests/test_linalg.py +++ /dev/null @@ -1,36 +0,0 @@ -import pytest - -import numpy as np -from numpy.testing import assert_allclose - -import finch - -arr1d = np.array([1, -1, 2, 3]) -arr2d = np.array([[1, 2, 0, 4, 0], [0, -2, 1, 0, 1]]) - - -@pytest.mark.parametrize("arr", [arr1d, arr2d]) -@pytest.mark.parametrize("keepdims", [True, False]) -@pytest.mark.parametrize( - "ord", - [ - 0, - 1, - 10, - finch.inf, - -finch.inf, - pytest.param( - 2, - marks=pytest.mark.skip( - reason="https://github.com/finch-tensor/Finch.jl/pull/709" - ), - ), - ], -) -def test_vector_norm(arr, keepdims, ord): - tns = finch.asarray(arr) - - actual = finch.linalg.vector_norm(tns, keepdims=keepdims, ord=ord) - expected = np.linalg.vector_norm(arr, keepdims=keepdims, ord=ord) - - assert_allclose(actual.todense(), expected) diff --git a/tests/test_ops.py b/tests/test_ops.py deleted file mode 100644 index b2dd340..0000000 --- a/tests/test_ops.py +++ /dev/null @@ -1,440 +0,0 @@ -from functools import reduce - -import pytest - -import numpy as np -from numpy.testing import assert_allclose, assert_equal - -import juliacall as jc - -import finch - -arr1d = np.array([1, 1, 2, 3]) -arr2d = np.array([[1, 2, 0, 0], [0, 1, 0, 1]]) -arr3d = np.array( - [ - [[0, 1, 0, 0], [1, 0, 0, 3]], - [[4, 0, -1, 0], [2, 2, 0, 0]], - [[0, 0, 0, 0], [1, 5, 0, 3]], - ] -) - - -@pytest.fixture( - scope="module", - params=[finch.DefaultScheduler(), finch.GalleyScheduler()], - ids=["default", "galley"], -) -def opt(request): - finch.set_optimizer(request.param) - yield request.param - - -def test_eager(arr3d, opt): - A_finch = finch.Tensor(arr3d) - B_finch = finch.Tensor(arr2d) - - result = finch.multiply(A_finch, B_finch) - - assert_equal(result.todense(), np.multiply(arr3d, arr2d)) - - -def test_lazy_mode(arr3d, opt): - A_finch = finch.Tensor(arr3d) - B_finch = finch.Tensor(arr2d) - C_finch = finch.Tensor(arr1d) - - @finch.compiled(opt=opt) - def my_custom_fun(arr1, arr2, arr3): - temp = finch.multiply(arr1, arr2) - temp = finch.divide(temp, arr3) - reduced = finch.sum(temp, axis=(0, 1)) - return finch.add(temp, reduced) - - result = my_custom_fun(A_finch, B_finch, C_finch) - - temp = np.divide(np.multiply(arr3d, arr2d), arr1d) - expected = np.add(temp, np.sum(temp, axis=(0, 1))) - assert_equal(result.todense(), expected) - - A_lazy = finch.lazy(A_finch) - B_lazy = finch.lazy(B_finch) - mul_lazy = finch.multiply(A_lazy, B_lazy) - result = finch.compute(mul_lazy) - - assert_equal(result.todense(), np.multiply(arr3d, arr2d)) - - -def test_lazy_mode_mult_output(opt): - A_finch = finch.Tensor(arr1d) - B_finch = finch.Tensor(arr2d) - - @finch.compiled(opt=opt) - def mult_out_fun(arr1, arr2): - out1 = finch.add(arr1, arr2) - out2 = finch.multiply(arr1, arr2) - out3 = arr2 ** finch.asarray(2) - return out1, out2, out3 - - res1, res2, res3 = mult_out_fun(A_finch, B_finch) - - assert_equal(res1.todense(), np.add(arr1d, arr2d)) - assert_equal(res2.todense(), np.multiply(arr1d, arr2d)) - assert_equal(res3.todense(), arr2d**2) - - -def test_lazy_mode_heterogenous_output(): - A_finch = finch.Tensor(arr1d) - B_finch = finch.Tensor(arr2d) - - @finch.compiled() - def heterogenous_fun(a: list[finch.Tensor], b: int): - sum_a = reduce(lambda x1, x2: x1 + x2, a) - b_squared = b**2 - return (a, sum_a, (b, "text"), {"key1": 12, "key2": b_squared}) - - ret = heterogenous_fun([A_finch, B_finch], 3) - - assert type(ret) is tuple - assert len(ret) == 4 - assert type(ret[0]) is list - assert len(ret[0]) == 2 - assert_equal(ret[0][0].todense(), arr1d) - assert_equal(ret[0][1].todense(), arr2d) - assert_equal(ret[1].todense(), arr1d + arr2d) - assert ret[2] == (3, "text") - assert type(ret[3]) is dict - assert ret[3] == {"key1": 12, "key2": 9} - - -@pytest.mark.parametrize( - "func_name", - [ - "log", - "log10", - "log1p", - "log2", - "sqrt", - "sign", - "round", - "exp", - "expm1", - "floor", - "ceil", - "isnan", - "isfinite", - "isinf", - "square", - "trunc", - ], -) -def test_elemwise_ops_1_arg(arr3d, func_name, opt): - arr = arr3d + 1.6 - A_finch = finch.Tensor(arr) - - actual = getattr(finch, func_name)(A_finch) - expected = getattr(np, func_name)(arr) - - assert_allclose(actual.todense(), expected) - - -@pytest.mark.parametrize("func_name", ["real", "imag", "conj"]) -@pytest.mark.parametrize("dtype", [np.complex128, np.complex64, np.float64, np.int64]) -def test_elemwise_complex_ops_1_arg(func_name, dtype, opt): - arr = np.asarray([[1 + 1j, 2 + 2j], [3 + 3j, 4 - 4j], [-5 - 5j, -6 - 6j]]).astype( - dtype - ) - arr_finch = finch.asarray(arr) - - actual = getattr(finch, func_name)(arr_finch) - expected = getattr(np, func_name)(arr) - - assert_allclose(actual.todense(), expected) - assert actual.todense().dtype == expected.dtype - - -@pytest.mark.parametrize( - "meth_name", - ["__pos__", "__neg__", "__abs__", "__invert__"], -) -def test_elemwise_tensor_ops_1_arg(arr3d, meth_name, opt): - A_finch = finch.Tensor(arr3d) - - actual = getattr(A_finch, meth_name)() - expected = getattr(arr3d, meth_name)() - - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize( - "func_name", - ["logaddexp", "logical_and", "logical_or", "logical_xor"], -) -def test_elemwise_ops_2_args(arr3d, func_name, opt): - arr2d = np.array([[0, 3, 2, 0], [0, 0, 3, 2]]) - if func_name.startswith("logical"): - arr3d = arr3d.astype(bool) - arr2d = arr2d.astype(bool) - A_finch = finch.Tensor(arr3d) - B_finch = finch.Tensor(arr2d) - - actual = getattr(finch, func_name)(A_finch, B_finch) - expected = getattr(np, func_name)(arr3d, arr2d) - - assert_allclose(actual.todense(), expected) - - -@pytest.mark.parametrize( - "meth_name", - [ - "__add__", - "__mul__", - "__sub__", - "__truediv__", - "__floordiv__", - "__mod__", - "__pow__", - "__and__", - "__or__", - "__xor__", - "__lshift__", - "__rshift__", - "__lt__", - "__le__", - "__gt__", - "__ge__", - "__eq__", - "__ne__", - ], -) -def test_elemwise_tensor_ops_2_args(arr3d, meth_name, opt): - arr2d = np.array([[2, 3, 2, 3], [3, 2, 3, 2]]) - A_finch = finch.Tensor(arr3d) - B_finch = finch.Tensor(arr2d) - - actual = getattr(A_finch, meth_name)(B_finch) - expected = getattr(arr3d, meth_name)(arr2d) - - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize( - "func_name", - [ - "sum", - "prod", - "max", - "min", - "any", - "all", - "mean", - pytest.param("std", marks=pytest.mark.xfail(reason="TODO: to debug")), - pytest.param("var", marks=pytest.mark.xfail(reason="TODO: to debug")), - ], -) -@pytest.mark.parametrize("axis", [None, -1, 1, (0, 1), (0, 1, 2)]) -def test_reductions(arr3d, func_name, axis, opt): - A_finch = finch.Tensor(arr3d) - - actual = getattr(finch, func_name)(A_finch, axis=axis) - expected = getattr(np, func_name)(arr3d, axis=axis) - - assert_equal(actual.todense(), expected) - - -@pytest.mark.skip(reason="TODO: to debug") -@pytest.mark.parametrize("func_name", ["argmax", "argmin"]) -@pytest.mark.parametrize("axis", [None, -1, 1, 2, (0, 1, 2)]) -def test_arg_reductions(arr3d, func_name, axis, opt): - A_finch = finch.Tensor(arr3d) - - actual = getattr(finch, func_name)(A_finch, axis=axis) - expected = getattr(np, func_name)(arr3d, axis=axis) - - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize("axis", [-1, 1, (0, 1), (0, 1, 2)]) -def test_expand_dims(arr3d, axis, opt): - A_finch = finch.Tensor(arr3d) - - actual = finch.expand_dims(A_finch, axis=axis) - expected = np.expand_dims(arr3d, axis=axis) - - assert_equal(actual.todense(), expected) - - actual = finch.squeeze(actual, axis=axis) - expected = np.squeeze(expected, axis=axis) - - assert_equal(actual.todense(), expected) - - -@pytest.mark.skip(reason="TODO: AssertionError") -@pytest.mark.parametrize("offset", [-1, 0, 1]) -def test_diagonal_2d_array(offset, opt): - arr2d = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) - A_finch = finch.Tensor(arr2d) - - actual = finch.diagonal(A_finch, offset=offset) - expected = np.diagonal(arr2d, offset=offset) - - assert_equal(actual.todense(), expected) - - -@pytest.mark.skip(reason="TODO: UndefVarError: `d_ijk` not defined in local scope") -@pytest.mark.parametrize("offset", [-1, 0, 1]) -def test_diagonal_high_dimensional_array(offset, opt): - rng = np.random.default_rng(42) - arr_high_dim = rng.random((4, 3, 5, 6)) - A_finch = finch.Tensor(arr_high_dim) - - actual = finch.diagonal(A_finch, offset=offset) - expected = np.diagonal(arr_high_dim, offset=offset) - - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize("func_name", ["sum", "prod"]) -@pytest.mark.parametrize("axis", [None, 0, 1]) -@pytest.mark.parametrize( - "in_dtype, dtype, expected_dtype", - [ - (finch.int64, None, np.int64), - (finch.int16, None, np.int64), - (finch.uint8, None, np.uint64), - (finch.int64, finch.float32, np.float32), - (finch.float64, finch.complex128, np.complex128), - ], -) -def test_sum_prod_dtype_arg( - arr3d, func_name, axis, in_dtype, dtype, expected_dtype, opt -): - arr_finch = finch.asarray(np.abs(arr3d), dtype=in_dtype) - - actual = getattr(finch, func_name)(arr_finch, axis=axis, dtype=dtype).todense() - - assert actual.dtype == expected_dtype - - -@pytest.mark.parametrize( - "storage", - [ - None, - ( - finch.Storage(finch.SparseList(finch.Element(np.int64(0))), order="C"), - finch.Storage( - finch.Dense(finch.SparseList(finch.Element(np.int64(0)))), order="C" - ), - finch.Storage( - finch.Dense( - finch.SparseList(finch.SparseList(finch.Element(np.int64(0)))) - ), - order="C", - ), - ), - ], -) -def test_tensordot(arr3d, storage, opt): - A_finch = finch.Tensor(arr1d) - B_finch = finch.Tensor(arr2d) - C_finch = finch.Tensor(arr3d) - if storage is not None: - A_finch = A_finch.to_storage(storage[0]) - B_finch = B_finch.to_storage(storage[1]) - C_finch = C_finch.to_storage(storage[2]) - - actual = finch.tensordot(B_finch, B_finch) - expected = np.tensordot(arr2d, arr2d) - assert_equal(actual.todense(), expected) - - actual = finch.tensordot(B_finch, B_finch, axes=(1, 1)) - expected = np.tensordot(arr2d, arr2d, axes=(1, 1)) - assert_equal(actual.todense(), expected) - - actual = finch.tensordot( - C_finch, finch.permute_dims(C_finch, (2, 1, 0)), axes=((2, 0), (0, 2)) - ) - expected = np.tensordot(arr3d, arr3d.T, axes=((2, 0), (0, 2))) - assert_equal(actual.todense(), expected) - - actual = finch.tensordot(C_finch, A_finch, axes=(2, 0)) - expected = np.tensordot(arr3d, arr1d, axes=(2, 0)) - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize( - ("a", "b"), - [ - (arr2d, arr2d.mT), - (arr2d, arr3d.mT), - (arr2d.mT, arr3d), - (arr3d, arr3d.mT), - (arr1d, arr1d), - (arr1d, arr2d.mT), - (arr2d, arr1d), - (arr1d, arr3d.mT), - (arr3d, arr1d), - ], -) -def test_matmul(opt, a: np.ndarray, b: np.ndarray): - A_finch = finch.Tensor(a) - B_finch = finch.Tensor(b) - - expected = a @ b - actual = A_finch @ B_finch - - assert_equal(actual.todense(), expected) - - if a.ndim >= 2 and b.ndim >= 2: - At_finch = A_finch.mT - Bt_finch = B_finch.mT - - assert_equal((Bt_finch @ At_finch).todense(), expected.mT) - - -def test_matmul_dimension_mismatch(opt): - A_finch = finch.Tensor(arr2d) - B_finch = finch.Tensor(arr3d) - - with pytest.raises(jc.JuliaError, match="DimensionMismatch"): - A_finch @ B_finch - - -def test_negative__mod__(opt): - arr = np.array([-1, 0, 0, -2, -3, 0]) - arr_finch = finch.asarray(arr) - - actual = arr_finch % 5 - expected = arr % 5 - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize("force_materialization", [False, True]) -def test_recursive_compiled( - opt, force_materialization: bool, arr3d: finch.Tensor -) -> None: - decorator = finch.compiled(opt=opt, force_materialization=force_materialization) - - @decorator - def my_custom_fun_inner( - arr1: finch.Tensor, arr2: finch.Tensor, arr3: finch.Tensor - ) -> finch.Tensor: - temp = finch.multiply(arr1, arr2) - temp = finch.divide(temp, arr3) - reduced = finch.sum(temp, axis=(0, 1)) - return finch.add(temp, reduced) - - @decorator - def my_custom_fun_outer( - arr1: finch.Tensor, arr2: finch.Tensor, arr3: finch.Tensor - ) -> finch.Tensor: - arr = my_custom_fun_inner(arr1, arr2, arr3) - assert arr.is_computed() == force_materialization - return arr - - A_finch = finch.Tensor(arr3d) - B_finch = finch.Tensor(arr2d) - C_finch = finch.Tensor(arr1d) - - result = my_custom_fun_outer(A_finch, B_finch, C_finch) - assert result.is_computed() diff --git a/tests/test_scipy_constructors.py b/tests/test_scipy_constructors.py deleted file mode 100644 index d472923..0000000 --- a/tests/test_scipy_constructors.py +++ /dev/null @@ -1,123 +0,0 @@ -import pytest - -import numpy as np -import scipy.sparse as sp -from numpy.testing import assert_equal - -import finch -from finch.tensor import _eq_scalars - - -def test_scipy_coo(arr2d): - sp_arr = sp.coo_matrix(arr2d, dtype=np.int64) - finch_arr = finch.Tensor(sp_arr) - lvl = finch_arr._obj.body.lvl - - assert np.shares_memory(sp_arr.row, lvl.tbl[1].data) - assert np.shares_memory(sp_arr.col, lvl.tbl[0].data) - assert np.shares_memory(sp_arr.data, lvl.lvl.val) - - assert_equal(finch_arr.todense(), sp_arr.todense()) - new_arr = finch.permute_dims(finch_arr, (1, 0)) - assert_equal(new_arr.todense(), sp_arr.todense().transpose()) - - -@pytest.mark.parametrize("cls", [sp.csc_matrix, sp.csr_matrix]) -def test_scipy_compressed2d(arr2d, cls): - sp_arr = cls(arr2d, dtype=np.int64) - finch_arr = finch.Tensor(sp_arr) - lvl = finch_arr._obj.body.lvl.lvl - - assert np.shares_memory(sp_arr.indices, lvl.idx.data) - assert np.shares_memory(sp_arr.indptr, lvl.ptr.data) - assert np.shares_memory(sp_arr.data, lvl.lvl.val) - - assert_equal(finch_arr.todense(), sp_arr.todense()) - new_arr = finch.permute_dims(finch_arr, (1, 0)) - assert_equal(new_arr.todense(), sp_arr.todense().transpose()) - - -@pytest.mark.parametrize( - "format_with_cls_with_order", - [ - ("coo", sp.coo_matrix, "C"), - ("coo", sp.coo_matrix, "F"), - ("csc", sp.csc_matrix, "F"), - ("csr", sp.csr_matrix, "C"), - ], -) -@pytest.mark.parametrize("fill_value_in", [0, finch.inf, finch.nan, 5, None]) -@pytest.mark.parametrize("fill_value_out", [0, finch.inf, finch.nan, 5, None]) -def test_to_scipy_sparse(format_with_cls_with_order, fill_value_in, fill_value_out): - format, sp_class, order = format_with_cls_with_order - np_arr = np.random.default_rng(0).random((4, 5)) - np_arr = np.array(np_arr, order=order) - - finch_arr = finch.asarray(np_arr, format=format, fill_value=fill_value_in) - - if not ( - fill_value_in in {0, None} and fill_value_out in {0, None} - ) and not _eq_scalars(fill_value_in, fill_value_out): - match_fill_value_out = 0 if fill_value_out is None else fill_value_out - with pytest.raises( - ValueError, - match=( - rf"Can only convert arrays with \[{match_fill_value_out}\] " - "fill-values to a Scipy sparse matrix." - ), - ): - finch_arr.to_scipy_sparse(accept_fv=fill_value_out) - return - - actual = finch_arr.to_scipy_sparse(accept_fv=fill_value_out) - - assert isinstance(actual, sp_class) - assert_equal(actual.todense(), np_arr) - - -def test_to_scipy_sparse_invalid_input(): - finch_arr = finch.asarray(np.ones((3, 3, 3)), format="dense") - - with pytest.raises(ValueError, match="Can only convert a 2-dimensional array"): - finch_arr.to_scipy_sparse() - - finch_arr = finch.asarray(np.ones((3, 4)), format="dense") - - with pytest.raises( - ValueError, match="Tensor can't be converted to scipy.sparse object" - ): - finch_arr.to_scipy_sparse() - - -@pytest.mark.parametrize( - "format_with_pattern", - [ - ("coo", "SparseCOO"), - ("csr", "SparseList"), - ("csc", "SparseList"), - ("bsr", "SparseCOO"), - ("dok", "SparseCOO"), - ], -) -@pytest.mark.parametrize("fill_value", [0, finch.inf, finch.nan, 5, None]) -def test_from_scipy_sparse(format_with_pattern, fill_value): - format, pattern = format_with_pattern - sp_arr = sp.random(10, 5, density=0.1, format=format) - - result = finch.Tensor.from_scipy_sparse(sp_arr, fill_value=fill_value) - assert pattern in str(result) - fill_value = 0 if fill_value is None else fill_value - assert _eq_scalars(result.fill_value, fill_value) - - -@pytest.mark.parametrize("format", ["coo", "bsr"]) -def test_non_canonical_format(format): - sp_arr = sp.random(3, 4, density=0.5, format=format) - - with pytest.raises( - ValueError, match="Unable to avoid copy while creating an array" - ): - finch.asarray(sp_arr, copy=False) - - finch_arr = finch.asarray(sp_arr) - assert_equal(finch_arr.todense(), sp_arr.toarray()) diff --git a/tests/test_sparse.py b/tests/test_sparse.py deleted file mode 100644 index 17c6318..0000000 --- a/tests/test_sparse.py +++ /dev/null @@ -1,424 +0,0 @@ -import pytest - -import numpy as np -from numpy.testing import assert_equal - -import sparse - -import finch - -parametrize_optimizer = pytest.mark.parametrize( - "opt", [finch.DefaultScheduler(), finch.GalleyScheduler()] -) - - -@pytest.mark.parametrize( - "dtype,jl_dtype", - [ - (np.int64, finch.int64), - (np.float64, finch.float64), - (np.complex128, finch.complex128), - ], -) -@pytest.mark.parametrize("order", ["C", "F", None]) -def test_wrappers(dtype, jl_dtype, order): - A = np.array([[0, 0, 4], [1, 0, 0], [2, 0, 5], [3, 0, 0]], dtype=dtype, order=order) - B = np.array(np.stack([A, A], axis=2, dtype=dtype), order=order) - - B_finch = finch.Tensor(B) - - storage = finch.Storage( - finch.Dense(finch.SparseList(finch.SparseList(finch.Element(dtype(0.0))))), - order=order, - ) - B_finch = B_finch.to_storage(storage) - - assert B_finch.shape == B.shape - assert B_finch.dtype == jl_dtype - assert_equal(B_finch.todense(), B) - - storage = finch.Storage( - finch.Dense(finch.Dense(finch.Element(dtype(1.0)))), order=order - ) - A_finch = finch.Tensor(A).to_storage(storage) - - assert A_finch.shape == A.shape - assert A_finch.dtype == jl_dtype - assert_equal(A_finch.todense(), A) - assert A_finch.todense().dtype == A.dtype and B_finch.todense().dtype == B.dtype - - -@pytest.mark.parametrize("dtype", [np.int64, np.float64, np.complex128]) -@pytest.mark.parametrize("order", ["C", "F", None]) -@pytest.mark.parametrize("copy", [True, False, None]) -def test_copy_fully_dense(dtype, order, copy, arr3d): - arr = np.array(arr3d, dtype=dtype, order=order) - arr_finch = finch.Tensor(arr, copy=copy) - arr_todense = arr_finch.todense() - - assert_equal(arr_todense, arr) - if copy: - assert not np.shares_memory(arr_todense, arr) - else: - assert np.shares_memory(arr_todense, arr) - - -def test_coo(rng): - coords = ( - np.asarray([0, 1, 2, 3, 4], dtype=np.intp), - np.asarray([0, 1, 2, 3, 4], dtype=np.intp), - ) - data = rng.random(5) - - arr_pydata = sparse.COO(np.vstack(coords), data, shape=(5, 5)) - arr = arr_pydata.todense() - arr_finch = finch.Tensor.construct_coo(coords, data, shape=(5, 5)) - - assert_equal(arr_finch.todense(), arr) - assert arr_finch.todense().dtype == data.dtype - - -@pytest.mark.parametrize( - "classes", - [ - (sparse._compressed.CSC, finch.Tensor.construct_csc), - (sparse._compressed.CSR, finch.Tensor.construct_csr), - ], -) -def test_compressed2d(rng, classes): - sparse_class, finch_class = classes - indices, indptr, data = np.arange(5), np.arange(6), rng.random(5) - - arr_pydata = sparse_class((data, indices, indptr), shape=(5, 5)) - arr = arr_pydata.todense() - arr_finch = finch_class((data, indices, indptr), shape=(5, 5)) - - assert_equal(arr_finch.todense(), arr) - assert arr_finch.todense().dtype == data.dtype - - -def test_csf(arr3d): - arr = arr3d - dtype = np.int64 - - data = np.array([4, 1, 2, 1, 1, 2, 5, -1, 3, 3], dtype=dtype) - indices_list = [ - np.array([1, 0, 1, 2, 0, 1, 2, 1, 0, 2], dtype=dtype), - np.array([0, 1, 0, 1, 0, 1], dtype=dtype), - ] - indptr_list = [ - np.array([0, 1, 4, 5, 7, 8, 10], dtype=dtype), - np.array([0, 2, 4, 5, 6], dtype=dtype), - ] - - arr_finch = finch.Tensor.construct_csf( - (data, indices_list, indptr_list), shape=(3, 2, 4) - ) - - assert_equal(arr_finch.todense(), arr) - assert arr_finch.todense().dtype == data.dtype - - -@pytest.mark.parametrize( - "permutation", - [(0, 1, 2), (2, 1, 0), (0, 2, 1), (1, 2, 0), (2, 0, 1), (1, 0, 2)], -) -@pytest.mark.parametrize( - "format", - [ - finch.Dense(finch.SparseList(finch.SparseList(finch.Element(0)))), - finch.Dense(finch.Dense(finch.Dense(finch.Element(0)))), - ], -) -@pytest.mark.parametrize("order", ["C", "F"]) -@parametrize_optimizer -def test_permute_dims(arr3d, permutation, format, order, opt): - finch.set_optimizer(opt) - arr = np.array(arr3d, order=order) - storage = finch.Storage(format, order=order) - - arr_finch = finch.Tensor(arr).to_storage(storage) - - actual_eager_mode = finch.permute_dims(arr_finch, permutation) - actual_lazy_mode = finch.compute( - finch.permute_dims(finch.lazy(arr_finch), permutation) - ) - expected = np.transpose(arr, permutation) - - assert_equal(actual_eager_mode.todense(), expected) - assert_equal(actual_lazy_mode.todense(), expected) - - actual_eager_mode = finch.permute_dims(actual_eager_mode, permutation) - actual_lazy_mode = finch.compute( - finch.permute_dims(finch.lazy(actual_lazy_mode), permutation) - ) - expected = np.transpose(expected, permutation) - - assert_equal(actual_eager_mode.todense(), expected) - assert_equal(actual_lazy_mode.todense(), expected) - - # test `.mT` - actual_eager_mode = arr_finch.mT - actual_lazy_mode = finch.compute(finch.lazy(arr_finch).mT) - expected = arr.mT - - assert_equal(actual_eager_mode.todense(), expected) - assert_equal(actual_lazy_mode.todense(), expected) - - -@pytest.mark.parametrize("src_dest", [(0, 1), (1, 0), (-1, 2), (-2, -1), (1, 1)]) -@parametrize_optimizer -def test_moveaxis(arr3d, src_dest, opt): - finch.set_optimizer(opt) - src, dest = src_dest - arr_finch = finch.Tensor(arr3d) - - actual = finch.moveaxis(arr_finch, src, dest) - expected = np.moveaxis(arr3d, src, dest) - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize("order", ["C", "F"]) -@parametrize_optimizer -def test_astype(arr3d, order, opt): - finch.set_optimizer(opt) - arr = np.array(arr3d, order=order, dtype=np.int64) - storage = finch.Storage( - finch.Dense(finch.SparseList(finch.SparseList(finch.Element(np.int64(0))))), - order=order, - ) - arr_finch = finch.Tensor(arr).to_storage(storage) - - result = finch.astype(arr_finch, finch.int64) - assert result is not arr_finch - result = result.todense() - assert_equal(result, arr) - assert result.dtype == arr.dtype - - result = finch.astype(arr_finch, finch.int64, copy=False) - assert result is arr_finch - result = result.todense() - assert_equal(result, arr) - assert result.dtype == arr.dtype - - result = finch.astype(arr_finch, finch.float32).todense() - arr = arr.astype(np.float32) - assert_equal(result, arr) - assert result.dtype == arr.dtype - - with pytest.raises( - ValueError, match="Unable to avoid a copy while casting in no-copy mode." - ): - finch.astype(arr_finch, finch.float64, copy=False) - - -@pytest.mark.parametrize("random_state", [42, np.random.default_rng(42)]) -@parametrize_optimizer -def test_random(random_state, opt): - finch.set_optimizer(opt) - result = finch.random((10, 20, 30), density=0.0, random_state=random_state) - expected = sparse.random((10, 20, 30), density=0.0, random_state=random_state) - - assert_equal(result.todense(), expected.todense()) - - # test reproducible runs - run1 = finch.random((20, 20), density=0.8, random_state=0) - run2 = finch.random((20, 20), density=0.8, random_state=0) - run3 = finch.random((20, 20), density=0.8, random_state=0) - assert_equal(run1.todense(), run2.todense()) - assert_equal(run1.todense(), run3.todense()) - - -@pytest.mark.parametrize("order", ["C", "F"]) -@pytest.mark.parametrize("format", ["coo", "csr", "csc", "csf", "dense", None]) -@parametrize_optimizer -def test_asarray(arr2d, arr3d, order, format, opt): - finch.set_optimizer(opt) - arr = arr3d if format == "csf" else arr2d - arr = np.array(arr, order=order) - arr_finch = finch.Tensor(arr) - - result = finch.asarray(arr_finch, format=format) - assert_equal(result.todense(), arr) - - -@pytest.mark.parametrize( - "arr,new_shape", - [ - (np.arange(10), (2, 5)), - (np.ones((10, 10)), (100,)), - (np.ones((3, 4, 5)), (5, 2, 2, 3)), - (np.arange(1), (1, 1, 1, 1)), - (np.arange(1).reshape((1, 1, 1)), (1,)), - (np.arange(1).reshape((1, 1)), ()), - (np.zeros((10, 1, 2)), (1, 5, 4, 1)), - (np.int64(0), ()), - (np.int64(0), (1, 1)), - ], -) -@pytest.mark.parametrize("order", ["C", "F"]) -@parametrize_optimizer -def test_reshape(arr, new_shape, order, opt): - finch.set_optimizer(opt) - - arr = np.array(arr, order=order) - arr_finch = finch.Tensor(arr) - - res = finch.reshape(arr_finch, new_shape) - assert_equal(res.todense(), arr.reshape(new_shape)) - - -@pytest.mark.parametrize("shape", [10, (3, 3), (2, 1, 5)]) -@pytest.mark.parametrize("dtype_name", [None, "int64", "float64"]) -@pytest.mark.parametrize("format", ["coo", "dense"]) -@parametrize_optimizer -def test_full_ones_zeros_empty(shape, dtype_name, format, opt): - finch.set_optimizer(opt) - - jl_dtype = getattr(finch, dtype_name) if dtype_name is not None else None - np_dtype = getattr(np, dtype_name) if dtype_name is not None else None - - res = finch.full(shape, 2.0, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.full(shape, 2.0, np_dtype)) - res = finch.full_like(res, 3.0, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.full(shape, 3.0, np_dtype)) - - res = finch.ones(shape, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.ones(shape, np_dtype)) - res = finch.ones_like(res, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.ones(shape, np_dtype)) - - res = finch.zeros(shape, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.zeros(shape, np_dtype)) - res = finch.zeros_like(res, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.zeros(shape, np_dtype)) - - res = finch.empty(shape, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.empty(shape, np_dtype)) - res = finch.empty_like(res, dtype=jl_dtype, format=format) - assert_equal(res.todense(), np.empty(shape, np_dtype)) - - -@pytest.mark.parametrize("func,arg", [(finch.asarray, np.zeros(3)), (finch.zeros, 3)]) -def test_device_keyword(func, arg): - func(arg, device="cpu") - - with pytest.raises( - ValueError, - match='Device not understood. Only "cpu" is allowed, but received: cuda', - ): - func(arg, device="cuda") - - -@pytest.mark.parametrize( - "order_and_format", - [("C", None), ("F", None), ("C", "coo"), ("F", "coo"), ("F", "csc")], -) -@parametrize_optimizer -def test_where(order_and_format, opt): - finch.set_optimizer(opt) - - order, format = order_and_format - cond = np.array( - [ - [True, False, False, False], - [False, True, True, False], - [True, False, True, True], - ], - order=order, - ) - arr1 = np.array([[0, 0, 0, 1], [0, 2, 0, 3], [1, 0, 0, 5]], order=order) - arr2 = np.array([10, 20, 30, 40], order=order) - - tns_cond = finch.asarray(cond, format=format) - arr1_cond = finch.asarray(arr1, format=format) - arr2_cond = finch.asarray(arr2) - - actual = finch.where(tns_cond, arr1_cond, arr2_cond) - expected = np.where(cond, arr1, arr2) - - assert_equal(actual.todense(), expected) - - -@pytest.mark.parametrize("order", ["C", "F"]) -@pytest.mark.parametrize( - "format_shape", - [ - ("coo", (80,)), - ("coo", (10, 5, 8)), - ("csf", (10, 5, 8)), - ("csr", (5, 10)), - ("csc", (5, 10)), - ], -) -@parametrize_optimizer -def test_nonzero(order, format_shape, opt): - finch.set_optimizer(opt) - - format, shape = format_shape - rng = np.random.default_rng(0) - arr = rng.random(shape) - arr = np.array(arr, order=order) - mask = arr < 0.8 - arr[mask] = 0.0 - - tns = finch.asarray(arr, format=format) - - actual = finch.nonzero(tns) - expected = np.nonzero(arr) - for actual_i, expected_i in zip(actual, expected, strict=False): - assert_equal(actual_i.todense(), expected_i) - - -@pytest.mark.parametrize("dtype_name", ["int64", "float64", "complex128"]) -@pytest.mark.parametrize("k", [0, -1, 1, -2, 2]) -@pytest.mark.parametrize("format", ["coo", "dense"]) -@parametrize_optimizer -def test_eye(dtype_name, k, format, opt): - finch.set_optimizer(opt) - - result = finch.eye(3, 4, k=k, dtype=getattr(finch, dtype_name), format=format) - expected = np.eye(3, 4, k=k, dtype=getattr(np, dtype_name)) - - assert_equal(result.todense(), expected) - - -@parametrize_optimizer -def test_to_scalar(opt): - finch.set_optimizer(opt) - - for obj, meth_name in [ - (True, "__bool__"), - (1, "__int__"), - (1.0, "__float__"), - (1, "__index__"), - (1 + 1j, "__complex__"), - ]: - tns = finch.asarray(np.asarray(obj)) - assert getattr(tns, meth_name)() == obj - - tns = finch.asarray(np.ones((2, 2))) - with pytest.raises( - ValueError, match=" can be computed for one-element tensors only." - ): - tns.__int__() - - -@pytest.mark.parametrize("dtype_name", [None, "int16", "float64"]) -@parametrize_optimizer -def test_arange_linspace(dtype_name, opt): - finch.set_optimizer(opt) - - if dtype_name is not None: - finch_dtype = getattr(finch, dtype_name) - np_dtype = getattr(np, dtype_name) - else: - finch_dtype = np_dtype = None - - result = finch.arange(10, 100, 5, dtype=finch_dtype) - expected = np.arange(10, 100, 5, dtype=np_dtype) - assert_equal(result.todense(), expected) - - result = finch.linspace(20, 80, 10, dtype=finch_dtype) - expected = np.linspace(20, 80, 10, dtype=np_dtype) - assert_equal(result.todense(), expected) From ad36ffbd0450b1a08448b73cf062084685d2d74d Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 20 Feb 2026 09:44:34 -0500 Subject: [PATCH 13/81] chore: cleanup --- src/finch/compiler.py | 39 +++++++++++++++++++++++++-------------- src/finch/typing.py | 8 -------- 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index cde8248..356c6cb 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,20 +1,19 @@ -from finch.tensor import FinchJLTensor +import operator -from finchlite.compile import NotationCompiler -from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary import finchlite.finch_notation.nodes as ntn -from finchlite.compile import dimension -from typing import Any +from finchlite.compile import NotationCompiler, dimension +from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary -import operator +from finch.tensor import FinchJLTensor -from .julia import jc, jl +from .julia import jl ops_map = {operator.add: "+", operator.mul: "*"} # https://github.com/finch-tensor/finch-tensor-lite/blob/main/tests/test_notation_interpreter.py + class FinchJLKernel(AssemblyKernel): def __init__(self, func_name, jl_code): # We store this code so that we can verify it in pytest @@ -26,6 +25,7 @@ def __call__(self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ... finch_fn = getattr(jl, self.func_name) return tuple(finch_fn(*[arg._obj for arg in args])) + class FinchJLLibrary(AssemblyLibrary): def __init__(self, kernel_dict): self.kernel_dict = kernel_dict @@ -69,12 +69,18 @@ def generate_julia(self, prgm, nestingLvl=0): return "" tab_str = " " * nestingLvl - return f"{tab_str}{self.generate_julia(lhs, nestingLvl)} = {self.generate_julia(rhs, nestingLvl)}" + return ( + f"{tab_str}{self.generate_julia(lhs, nestingLvl)} = " + f"{self.generate_julia(rhs, nestingLvl)}" + ) - case ntn.Declare(tns, init, op, shape): + case ntn.Declare(tns, init, op, _): # TODO: what is the purpose of op here tab_str = " " * nestingLvl - return f"{tab_str}@finch {self.generate_julia(tns, nestingLvl)} .= {self.generate_julia(init, nestingLvl)}" + return ( + f"{tab_str}@finch {self.generate_julia(tns, nestingLvl)} .= " + f"{self.generate_julia(init, nestingLvl)}" + ) case ntn.Return(val): tab_str = " " * nestingLvl @@ -94,9 +100,11 @@ def generate_julia(self, prgm, nestingLvl=0): if not is_outermost_loop: return f"{tab_str}for {idx.name} = _\n{loop_body}{tab_str}end\n" - else: - self.in_finch_block = False - return f"{tab_str}@finch begin\n{tab_str_1}for {idx.name} = _\n{loop_body}{tab_str_1}end\n{tab_str}end" + self.in_finch_block = False + return ( + f"{tab_str}@finch begin\n{tab_str_1}for {idx.name} = " + f"_\n{loop_body}{tab_str_1}end\n{tab_str}end" + ) case ntn.Access(tns, _, idxs): tns_str = self.generate_julia(tns, nestingLvl) @@ -121,7 +129,10 @@ def generate_julia(self, prgm, nestingLvl=0): cond_str = self.generate_julia(cond, nestingLvl) then_body_str = self.generate_julia(then_body, nestingLvl + 1) else_body_str = self.generate_julia(else_body, nestingLvl + 1) - return f"{tab_str}if {cond_str}\n{then_body_str}\n{tab_str}else\n{else_body_str}\n{tab_str}end" + return ( + f"{tab_str}if {cond_str}\n{then_body_str}\n" + f"{tab_str}else\n{else_body_str}\n{tab_str}end" + ) case ntn.Increment(lhs, rhs): tab_str = " " * nestingLvl diff --git a/src/finch/typing.py b/src/finch/typing.py index 125086e..a5bacb6 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -1,13 +1,5 @@ -from typing import Any, Literal - import juliacall as jc -spmatrix = Any - JuliaObj = jc.AnyValue -DType = jc.AnyValue # represents jl.DataType - -Device = Literal["cpu"] | None - number = int | float | bool | complex From 20b9a13cf5f41dc3481d6fd6e43d71b185144982 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 20 Feb 2026 10:58:59 -0500 Subject: [PATCH 14/81] feat: added scheduler and improved code organization --- pixi.toml | 3 -- pyproject.toml | 4 --- src/finch/compiler.py | 8 ++---- src/finch/scheduler.py | 40 ++++++++++++++++++++++++++ src/finch/tensor.py | 4 +-- tests/test_compiler.py | 64 +++++++++++++++++++++++++----------------- 6 files changed, 84 insertions(+), 39 deletions(-) create mode 100644 src/finch/scheduler.py diff --git a/pixi.toml b/pixi.toml index 25c2efa..0ae882a 100644 --- a/pixi.toml +++ b/pixi.toml @@ -22,9 +22,6 @@ numpy = ">=1.19" [feature.test.pypi-dependencies] pytest = "*" pytest-cov = "*" -sparse = ">=0.16,<0.17" -numba = ">=0.61" -scipy = "*" numpy = "==2.*" pytest-xdist = ">=3.6.1,<4" diff --git a/pyproject.toml b/pyproject.toml index e0ec2fe..59a86d3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,10 +8,8 @@ requires-python = ">=3.11,<3.14" dependencies = [ "juliapkg (>=0.1.16,<0.2.0)", "numpy (>=1.19,<2.4)", - "numba>=0.61,<0.63.1", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", - "finch-tensor-lite==0.2.0" ] [tool.poetry] @@ -23,8 +21,6 @@ test = [ "pytest>=7.4.4,<7.6", "pre-commit>=3.6.0,<3.9", "pytest-cov>=4.1.0,<4.2", - "sparse>=0.17.0,<0.17.1", - "scipy>=1.7,<1.17", ] [build-system] diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 356c6cb..e9985fa 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -4,16 +4,12 @@ from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary -from finch.tensor import FinchJLTensor - from .julia import jl +from .tensor import FinchJLTensor ops_map = {operator.add: "+", operator.mul: "*"} -# https://github.com/finch-tensor/finch-tensor-lite/blob/main/tests/test_notation_interpreter.py - - class FinchJLKernel(AssemblyKernel): def __init__(self, func_name, jl_code): # We store this code so that we can verify it in pytest @@ -34,6 +30,8 @@ def __getattr__(self, name: str) -> FinchJLKernel: return self.kernel_dict[name] +# Test with +# https://github.com/finch-tensor/finch-tensor-lite/blob/main/tests/test_notation_interpreter.py class FinchJLGenerator: def __init__(self): self.pack_dict = {} diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py new file mode 100644 index 0000000..3dd24b8 --- /dev/null +++ b/src/finch/scheduler.py @@ -0,0 +1,40 @@ +from typing import Any + +from finchlite.autoschedule import ( + DefaultLogicOptimizer, + LogicCompiler, + LogicExecutor, + LogicFormatter, + LogicNormalizer, + LogicStandardizer, +) +from finchlite.finch_logic import LogicLoader +from finchlite.interface.fuse import set_default_scheduler + +from .compiler import FinchJLCompiler +from .levels import Dense, Element + + +class FinchJLLogicFormatter(LogicFormatter): + def __init__( + self, + loader: LogicLoader | None = None, + ): + super().__init__(loader) + + def get_output_tns_ftype(self, fill_value: Any, shape_type: tuple[Any, ...]): + lvl = Element(fill_value) + for _ in len(shape_type): + lvl = Dense(lvl) + return lvl + + +COMPILE_JULIA = LogicNormalizer( + LogicExecutor( + DefaultLogicOptimizer( + LogicStandardizer(FinchJLLogicFormatter(LogicCompiler(FinchJLCompiler()))) + ) + ) +) + +set_default_scheduler(COMPILE_JULIA) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index d700f9a..b4ec00e 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -20,8 +20,8 @@ def __str__(self): class FinchJLTensor(_Display, EagerTensor): - def __init__(self, obj: jl.Finch.Tensor): - if isinstance(obj, jl.Finch.Tensor): + def __init__(self, obj: JuliaObj): + if isinstance(obj, JuliaObj): self._obj = obj else: raise ValueError(f"Raw julia object expected. Found: {type(obj)}") diff --git a/tests/test_compiler.py b/tests/test_compiler.py index f26585c..6c6b136 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1,39 +1,38 @@ +import operator + import pytest + import numpy as np + +from finchlite.compile import ExtentFType, dimension from finchlite.finch_notation.nodes import ( - Module, - Function, - Variable, - Block, + Access, Assign, + Block, Call, + Declare, + Freeze, + Function, + Increment, Literal, - Update, - Unpack, - Slot, Loop, - Increment, - Access, - Unwrap, - Freeze, + Module, Read, Repack, Return, - Declare, + Slot, + Unpack, + Unwrap, + Update, + Variable, ) -import operator -from finchlite import ftype -from finchlite.algebra import overwrite, promote_min -from finchlite.compile import ExtentFType, dimension -from finchlite.codegen import NumpyBuffer - from finch.compiler import FinchJLCompiler, FinchJLKernel +from finch.julia import jl +from finch.levels import Dense, Element from finch.tensor import FinchJLTensor -# Dummy data to obtain the bufferized ND array type -a = np.zeros(dtype=np.float64, shape=(3, 3)) -a_format = ftype(FinchJLTensor(a)) +a_format = Dense(Dense(Element(0))) @pytest.mark.skip @@ -207,13 +206,28 @@ def test_finchjl_compiler(finch_ntn: Module, julia_code): return C end""", ( - FinchJLTensor(np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]])), - FinchJLTensor(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])), - FinchJLTensor(np.array([[10, 11, 12], [13, 14, 15], [16, 17, 18]])), + FinchJLTensor( + jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0))), (3, 3)) + ), + FinchJLTensor( + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(0))), + [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + ) + ), + FinchJLTensor( + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(0))), + [[10, 11, 12], [13, 14, 15], [16, 17, 18]], + ) + ), ), ( FinchJLTensor( - np.array([[84, 90, 96], [201, 216, 231], [318, 342, 366]]) + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(0))), + [[84, 90, 96], [201, 216, 231], [318, 342, 366]], + ) ), ), ) From 2543815e321b0124291afc412b04b948d582f9bc Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Fri, 20 Feb 2026 11:10:11 -0500 Subject: [PATCH 15/81] merged tensor and levels file --- src/finch/scheduler.py | 2 +- src/finch/tensor.py | 114 ++++++++++++++++++++++++++++++++++++----- tests/test_compiler.py | 3 +- 3 files changed, 102 insertions(+), 17 deletions(-) diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py index 3dd24b8..38ae6a7 100644 --- a/src/finch/scheduler.py +++ b/src/finch/scheduler.py @@ -12,7 +12,7 @@ from finchlite.interface.fuse import set_default_scheduler from .compiler import FinchJLCompiler -from .levels import Dense, Element +from .tensor import Dense, Element class FinchJLLogicFormatter(LogicFormatter): diff --git a/src/finch/tensor.py b/src/finch/tensor.py index b4ec00e..ab7a121 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,37 +1,123 @@ -from finchlite import EagerTensor +from abc import abstractmethod +from typing import Any + +import numpy as np + +from finchlite import EagerTensor, Tensor, TensorFType from .julia import jl -from .typing import JuliaObj +from .typing import JuliaObj, number -# Singleton classes for levels types -# finch tensor lite, formatter stage -# level ftype without the need to create tthe object -# https://github.com/finch-tensor/finch-tensor-lite/blob/main/src/finchlite/autoschedule/formatter.py +# Abstract FTypes +class LevelFType(TensorFType): + def from_numpy(self, _) -> Tensor: + raise NotImplementedError -class _Display: - _obj: JuliaObj + def shape_type(self) -> tuple[type, ...]: + return tuple(self.element_type for _ in range(self.ndim)) - def __repr__(self): - return jl.sprint(jl.show, self._obj) - def __str__(self): - return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) +class NestedLevelFType(LevelFType): + def __init__(self, lvl: LevelFType): + self.lvl = lvl + + def ndims(self) -> np.intp: + return self.lvl.ndims + np.intp(1) + + def fill_value(self) -> Any: + return self.lvl.fill_value + + def element_type(self) -> Any: + return self.lvl.element_type + + def __call__(self, shape: tuple) -> Tensor: + return FinchJLTensor(jl.Finch.Tensor(self.create_jl_obj(), shape)) + + def __eq__(self, other): + return type(other) is type(self) and self.lvl == other.lvl + + def __hash__(self): + return hash((self.__class__.__name__, self.lvl.__hash__)) + + @abstractmethod + def create_jl_obj(self) -> JuliaObj: ... + + +# Concrete FTypes +class Element(LevelFType): + def __init__(self, fill_value: number): + self._fill_value = fill_value + + def ndims(self) -> np.intp: + return np.intp(0) + + def fill_value(self) -> Any: + return self._fill_value + + def element_type(self) -> Any: + return type(self._fill_value) + + def __call__(self, _) -> Tensor: + raise Exception("Cannot create an object of element type!") + def __eq__(self, other): + return isinstance(other, Element) and self._fill_value == other.fill_value -class FinchJLTensor(_Display, EagerTensor): + def __hash__(self): + return hash((self.__class__.__name__, self._fill_value)) + + def create_jl_obj(self) -> JuliaObj: + return jl.Element(self._fill_value) + + +class Dense(NestedLevelFType): + def create_jl_obj(self) -> JuliaObj: + return jl.Dense(self.lvl.create_jl_obj()) + + +class SparseList(NestedLevelFType): + def create_jl_obj(self) -> JuliaObj: + return jl.SparseList(self.lvl.create_jl_obj()) + + +class SparseByteMap(NestedLevelFType): + def create_jl_obj(self) -> JuliaObj: + return jl.SparseByteMap(self.lvl.create_jl_obj()) + + +# Tensor Class +class FinchJLTensor(EagerTensor): def __init__(self, obj: JuliaObj): if isinstance(obj, JuliaObj): self._obj = obj else: raise ValueError(f"Raw julia object expected. Found: {type(obj)}") - # TODO: figure out a way to walk through the levels and return the ftype + # TODO: figure out a way to walk through the levels and construct the ftype + self._ftype = Dense(Dense(Element(0))) + @property def ftype(self): """Returns the ftype of the buffer""" + return self._ftype @property def shape(self) -> tuple: """Shape of the tensor.""" return self.obj.shape + + def ndims(self) -> np.intp: + return self._ftype.ndims + + def fill_value(self) -> Any: + return self._ftype.fill_value + + def element_type(self) -> Any: + return self._ftype.element_type + + def __repr__(self): + return jl.sprint(jl.show, self._obj) + + def __str__(self): + return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 6c6b136..f08a575 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -29,8 +29,7 @@ from finch.compiler import FinchJLCompiler, FinchJLKernel from finch.julia import jl -from finch.levels import Dense, Element -from finch.tensor import FinchJLTensor +from finch.tensor import Dense, Element, FinchJLTensor a_format = Dense(Dense(Element(0))) From a2ab46d139cdc590355d176ace74c2eb2ffd4e78 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Mon, 23 Feb 2026 14:20:08 -0500 Subject: [PATCH 16/81] fix: resolved issues with tests --- src/finch/compiler.py | 8 +++++++- src/finch/tensor.py | 15 ++++++++------- tests/test_compiler.py | 28 ++++------------------------ 3 files changed, 19 insertions(+), 32 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index e9985fa..d9e3d9e 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -19,7 +19,13 @@ def __init__(self, func_name, jl_code): def __call__(self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...]: finch_fn = getattr(jl, self.func_name) - return tuple(finch_fn(*[arg._obj for arg in args])) + result = finch_fn(*[arg._obj for arg in args]) + + # The finch function returns tuples when multiple values are returned + # or a non-tuple when a single value is returned. + if not isinstance(result, tuple): + result = (result,) + return tuple(FinchJLTensor(res) for res in result) class FinchJLLibrary(AssemblyLibrary): diff --git a/src/finch/tensor.py b/src/finch/tensor.py index ab7a121..0079559 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -22,8 +22,8 @@ class NestedLevelFType(LevelFType): def __init__(self, lvl: LevelFType): self.lvl = lvl - def ndims(self) -> np.intp: - return self.lvl.ndims + np.intp(1) + def ndim(self) -> np.intp: + return self.lvl.ndim + np.intp(1) def fill_value(self) -> Any: return self.lvl.fill_value @@ -49,7 +49,7 @@ class Element(LevelFType): def __init__(self, fill_value: number): self._fill_value = fill_value - def ndims(self) -> np.intp: + def ndim(self) -> np.intp: return np.intp(0) def fill_value(self) -> Any: @@ -97,18 +97,16 @@ def __init__(self, obj: JuliaObj): # TODO: figure out a way to walk through the levels and construct the ftype self._ftype = Dense(Dense(Element(0))) - @property def ftype(self): """Returns the ftype of the buffer""" return self._ftype - @property def shape(self) -> tuple: """Shape of the tensor.""" return self.obj.shape - def ndims(self) -> np.intp: - return self._ftype.ndims + def ndim(self) -> np.intp: + return self._ftype.ndim def fill_value(self) -> Any: return self._ftype.fill_value @@ -116,6 +114,9 @@ def fill_value(self) -> Any: def element_type(self) -> Any: return self._ftype.element_type + def __eq__(self, other): + return isinstance(other, FinchJLTensor) and self._obj == other._obj + def __repr__(self): return jl.sprint(jl.show, self._obj) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index f08a575..39b853d 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -34,7 +34,6 @@ a_format = Dense(Dense(Element(0))) -@pytest.mark.skip @pytest.mark.parametrize( "finch_ntn, julia_code", [ @@ -205,30 +204,11 @@ def test_finchjl_compiler(finch_ntn: Module, julia_code): return C end""", ( - FinchJLTensor( - jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0))), (3, 3)) - ), - FinchJLTensor( - jl.Finch.Tensor( - jl.Dense(jl.Dense(jl.Element(0))), - [[1, 2, 3], [4, 5, 6], [7, 8, 9]], - ) - ), - FinchJLTensor( - jl.Finch.Tensor( - jl.Dense(jl.Dense(jl.Element(0))), - [[10, 11, 12], [13, 14, 15], [16, 17, 18]], - ) - ), - ), - ( - FinchJLTensor( - jl.Finch.Tensor( - jl.Dense(jl.Dense(jl.Element(0))), - [[84, 90, 96], [201, 216, 231], [318, 342, 366]], - ) - ), + FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0))), 3, 3)), + FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(1))), 3, 3)), + FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(2))), 3, 3)), ), + (FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(6))), 3, 3)),), ) ], ) From 117d223e966797e8787bb465eb52b1f08ea0a8fa Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Mon, 23 Feb 2026 15:30:29 -0500 Subject: [PATCH 17/81] feat: added support for array based constructor --- src/finch/scheduler.py | 3 --- tests/test_compiler.py | 31 +++++++++++++++++++++++++++---- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py index 38ae6a7..36b0289 100644 --- a/src/finch/scheduler.py +++ b/src/finch/scheduler.py @@ -9,7 +9,6 @@ LogicStandardizer, ) from finchlite.finch_logic import LogicLoader -from finchlite.interface.fuse import set_default_scheduler from .compiler import FinchJLCompiler from .tensor import Dense, Element @@ -36,5 +35,3 @@ def get_output_tns_ftype(self, fill_value: Any, shape_type: tuple[Any, ...]): ) ) ) - -set_default_scheduler(COMPILE_JULIA) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 39b853d..3517ba8 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -34,6 +34,7 @@ a_format = Dense(Dense(Element(0))) +@pytest.mark.skip @pytest.mark.parametrize( "finch_ntn, julia_code", [ @@ -204,11 +205,33 @@ def test_finchjl_compiler(finch_ntn: Module, julia_code): return C end""", ( - FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0))), 3, 3)), - FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(1))), 3, 3)), - FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(2))), 3, 3)), + FinchJLTensor( + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(0))), + np.array([[0, 0, 0], [0, 0, 0], [0, 0, 0]]), + ) + ), + FinchJLTensor( + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(1))), + np.array([[1, 1, 1], [1, 1, 1], [1, 1, 1]]), + ) + ), + FinchJLTensor( + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(2))), + np.array([[2, 2, 2], [2, 2, 2], [2, 2, 2]]), + ) + ), + ), + ( + FinchJLTensor( + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(0))), + np.array([[6, 6, 6], [6, 6, 6], [6, 6, 6]]), + ) + ), ), - (FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(6))), 3, 3)),), ) ], ) From b6de1217e36aaf54942ee66b25483d849508c69e Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Mon, 23 Feb 2026 15:49:56 -0500 Subject: [PATCH 18/81] feat: added proxy tesnor ftype --- src/finch/tensor.py | 48 ++++++++++++++++++++++++++++++++------------- 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 0079559..29ded0d 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -86,7 +86,38 @@ def create_jl_obj(self) -> JuliaObj: return jl.SparseByteMap(self.lvl.create_jl_obj()) -# Tensor Class +# Tensor Class and associated ftype +class FinchJLTensorFType(TensorFType): + def __init__(self, lvl): + self._lvl: NestedLevelFType = lvl + + def ndim(self) -> np.intp: + return self._lvl.ndim + + def fill_value(self) -> Any: + return self._lvl.fill_value + + def element_type(self) -> Any: + return self._lvl.element_type + + def shape_type(self) -> tuple[type, ...]: + return self._lvl.shape_type + + def __call__(self, shape: tuple) -> Tensor: + return self._lvl(shape) + + def from_numpy(self, _) -> Tensor: + raise NotImplementedError + + def __eq__(self, other): + if not isinstance(other, FinchJLTensorFType): + return False + return self._lvl == other._lvl + + def __hash__(self): + return hash(("FinchJLTensorFType", self._lvl)) + + class FinchJLTensor(EagerTensor): def __init__(self, obj: JuliaObj): if isinstance(obj, JuliaObj): @@ -94,26 +125,15 @@ def __init__(self, obj: JuliaObj): else: raise ValueError(f"Raw julia object expected. Found: {type(obj)}") - # TODO: figure out a way to walk through the levels and construct the ftype - self._ftype = Dense(Dense(Element(0))) - def ftype(self): """Returns the ftype of the buffer""" - return self._ftype + # TODO: figure out a way to walk through the levels and construct the ftype + return FinchJLTensorFType(Dense(Dense(Element(0)))) def shape(self) -> tuple: """Shape of the tensor.""" return self.obj.shape - def ndim(self) -> np.intp: - return self._ftype.ndim - - def fill_value(self) -> Any: - return self._ftype.fill_value - - def element_type(self) -> Any: - return self._ftype.element_type - def __eq__(self, other): return isinstance(other, FinchJLTensor) and self._obj == other._obj From 459452659087e36cec511d378caa604ebe513875 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Mon, 23 Feb 2026 17:48:33 -0500 Subject: [PATCH 19/81] chore: reorganizing to levels folder --- src/finch/__init__.py | 16 +++++++ src/finch/levels.py | 56 ++++++++++++------------- src/finch/tensor.py | 98 ++++++++----------------------------------- src/finch/utils.py | 81 +++++++++++++++++++++++++++++++++++ 4 files changed, 142 insertions(+), 109 deletions(-) create mode 100644 src/finch/utils.py diff --git a/src/finch/__init__.py b/src/finch/__init__.py index e69de29..6ef8cdd 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -0,0 +1,16 @@ +from .levels import Dense, Element, SparseByteMap, SparseList +from .scheduler import COMPILE_JULIA +from .tensor import ( + FinchJLTensor, + FinchJLTensorFType, +) + +__all__ = [ + "COMPILE_JULIA", + "Dense", + "Element", + "FinchJLTensor", + "FinchJLTensorFType", + "SparseByteMap", + "SparseList", +] diff --git a/src/finch/levels.py b/src/finch/levels.py index 284f6ae..3a0ddfa 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -6,10 +6,10 @@ from finchlite import Tensor, TensorFType from .julia import jl -from .tensor import FinchJLTensor from .typing import JuliaObj, number +# Abstract FTypes class LevelFType(TensorFType): def from_numpy(self, _) -> Tensor: raise NotImplementedError @@ -17,57 +17,55 @@ def from_numpy(self, _) -> Tensor: def shape_type(self) -> tuple[type, ...]: return tuple(self.element_type for _ in range(self.ndim)) + def __call__(self, _) -> Tensor: + raise Exception("Cannot create an object of this type!") -class Element(LevelFType): - def __init__(self, fill_value: number): - self._fill_value = fill_value - def ndims(self) -> np.intp: - return np.intp(0) +class NestedLevelFType(LevelFType): + def __init__(self, lvl: LevelFType): + self.lvl = lvl + + def ndim(self) -> np.intp: + return self.lvl.ndim + np.intp(1) def fill_value(self) -> Any: - return self._fill_value + return self.lvl.fill_value def element_type(self) -> Any: - return type(self._fill_value) - - def __call__(self, _) -> Tensor: - raise Exception("Cannot create an object of element type!") + return self.lvl.element_type def __eq__(self, other): - return isinstance(other, Element) and self._fill_value == other.fill_value + return type(other) is type(self) and self.lvl == other.lvl def __hash__(self): - return hash((self.__class__.__name__, self._fill_value)) + return hash((self.__class__.__name__, self.lvl.__hash__)) - def create_jl_obj(self) -> JuliaObj: - return jl.Element(self._fill_value) + @abstractmethod + def create_jl_obj(self) -> JuliaObj: ... -class NestedLevelFType(LevelFType): - def __init__(self, lvl: LevelFType): - self.lvl = lvl +# Concrete FTypes +class Element(LevelFType): + def __init__(self, fill_value: number): + self._fill_value = fill_value - def ndims(self) -> np.intp: - return self.lvl.ndims + np.intp(1) + def ndim(self) -> np.intp: + return np.intp(0) def fill_value(self) -> Any: - return self.lvl.fill_value + return self._fill_value def element_type(self) -> Any: - return self.lvl.element_type - - def __call__(self, shape: tuple) -> FinchJLTensor: - return FinchJLTensor(jl.Finch.Tensor(self.create_jl_obj(), shape)) + return type(self._fill_value) def __eq__(self, other): - return type(other) is type(self) and self.lvl == other.lvl + return isinstance(other, Element) and self._fill_value == other.fill_value def __hash__(self): - return hash((self.__class__.__name__, self.lvl.__hash__)) + return hash((self.__class__.__name__, self._fill_value)) - @abstractmethod - def create_jl_obj(self) -> JuliaObj: ... + def create_jl_obj(self) -> JuliaObj: + return jl.Element(self._fill_value) class Dense(NestedLevelFType): diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 29ded0d..2a178f4 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,4 +1,3 @@ -from abc import abstractmethod from typing import Any import numpy as np @@ -6,84 +5,9 @@ from finchlite import EagerTensor, Tensor, TensorFType from .julia import jl -from .typing import JuliaObj, number - - -# Abstract FTypes -class LevelFType(TensorFType): - def from_numpy(self, _) -> Tensor: - raise NotImplementedError - - def shape_type(self) -> tuple[type, ...]: - return tuple(self.element_type for _ in range(self.ndim)) - - -class NestedLevelFType(LevelFType): - def __init__(self, lvl: LevelFType): - self.lvl = lvl - - def ndim(self) -> np.intp: - return self.lvl.ndim + np.intp(1) - - def fill_value(self) -> Any: - return self.lvl.fill_value - - def element_type(self) -> Any: - return self.lvl.element_type - - def __call__(self, shape: tuple) -> Tensor: - return FinchJLTensor(jl.Finch.Tensor(self.create_jl_obj(), shape)) - - def __eq__(self, other): - return type(other) is type(self) and self.lvl == other.lvl - - def __hash__(self): - return hash((self.__class__.__name__, self.lvl.__hash__)) - - @abstractmethod - def create_jl_obj(self) -> JuliaObj: ... - - -# Concrete FTypes -class Element(LevelFType): - def __init__(self, fill_value: number): - self._fill_value = fill_value - - def ndim(self) -> np.intp: - return np.intp(0) - - def fill_value(self) -> Any: - return self._fill_value - - def element_type(self) -> Any: - return type(self._fill_value) - - def __call__(self, _) -> Tensor: - raise Exception("Cannot create an object of element type!") - - def __eq__(self, other): - return isinstance(other, Element) and self._fill_value == other.fill_value - - def __hash__(self): - return hash((self.__class__.__name__, self._fill_value)) - - def create_jl_obj(self) -> JuliaObj: - return jl.Element(self._fill_value) - - -class Dense(NestedLevelFType): - def create_jl_obj(self) -> JuliaObj: - return jl.Dense(self.lvl.create_jl_obj()) - - -class SparseList(NestedLevelFType): - def create_jl_obj(self) -> JuliaObj: - return jl.SparseList(self.lvl.create_jl_obj()) - - -class SparseByteMap(NestedLevelFType): - def create_jl_obj(self) -> JuliaObj: - return jl.SparseByteMap(self.lvl.create_jl_obj()) +from .levels import Dense, Element, NestedLevelFType +from .typing import JuliaObj +from .utils import add_missing_dims, add_plus_one, expand_ellipsis # Tensor Class and associated ftype @@ -104,7 +28,7 @@ def shape_type(self) -> tuple[type, ...]: return self._lvl.shape_type def __call__(self, shape: tuple) -> Tensor: - return self._lvl(shape) + return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), shape)) def from_numpy(self, _) -> Tensor: raise NotImplementedError @@ -134,6 +58,20 @@ def shape(self) -> tuple: """Shape of the tensor.""" return self.obj.shape + def __getitem__(self, key): + if not isinstance(key, tuple): + key = (key,) + + # standard indexing mode + key = expand_ellipsis(key, self.shape) + key = add_missing_dims(key, self.shape) + key = add_plus_one(key, self.shape) + + result = self._obj[key] + if jl.isa(result, jl.Finch.Tensor): + return FinchJLTensor(result) + return result + def __eq__(self, other): return isinstance(other, FinchJLTensor) and self._obj == other._obj diff --git a/src/finch/utils.py b/src/finch/utils.py new file mode 100644 index 0000000..accfa47 --- /dev/null +++ b/src/finch/utils.py @@ -0,0 +1,81 @@ +# Helper functions for indexing support + +import builtins + +import numpy as np +from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple + +from .julia import jl + + +def expand_ellipsis(key: tuple, shape: tuple[int, ...]) -> tuple: + ellipsis_pos = None + key_without_ellipsis = [] + # first we need to find the ellipsis and confirm it's the only one + for pos, idx in enumerate(key): + if idx is Ellipsis: + if ellipsis_pos is None: + ellipsis_pos = pos + else: + raise IndexError("an index can only have a single ellipsis ('...')") + else: + key_without_ellipsis.append(idx) + key = key_without_ellipsis + + # then we expand ellipsis with a full range + if ellipsis_pos is not None: + n_missing_idxs = len(shape) - builtins.sum(1 for k in key if k is not None) + key = key[:ellipsis_pos] + [slice(None)] * n_missing_idxs + key[ellipsis_pos:] + + return tuple(key) + + +def add_missing_dims(key: tuple, shape: tuple[int, ...]) -> tuple: + missing_dims = len(shape) - builtins.sum(1 for k in key if k is not None) + return key + (slice(None),) * missing_dims + + +def _slice_plus_one(s: slice, size: int) -> range: + step = s.step if s.step is not None else 1 + start_default = size if step < 0 else 1 + stop_default = 1 if step < 0 else size + + if s.start is not None: + start = normalize_axis_index(s.start, size) + 1 if s.start < size else size + else: + start = start_default + + if s.stop is not None: + stop_offset = 2 if step < 0 else 0 + stop = ( + normalize_axis_index(s.stop, size) + stop_offset if s.stop < size else size + ) + else: + stop = stop_default + + if (start, stop, step) == (1, size, 1): + return jl.Colon() + + return jl.range(start=start, step=step, stop=stop) + + +def add_plus_one(key: tuple, shape: tuple[int, ...]) -> tuple: + new_key = [] + sizes = iter(shape) + for idx in key: + if idx is None: + new_key.append(jl.nothing) + continue + + size = next(sizes) + if isinstance(idx, int): + new_key.append(normalize_axis_index(idx, size) + 1) + elif isinstance(idx, slice): + new_key.append(_slice_plus_one(idx, size)) + elif isinstance(idx, list | np.ndarray | tuple): + idx = normalize_axis_tuple(idx, size) + new_key.append(jl.Vector([i + 1 for i in idx])) + else: + new_key.append(idx) + + return tuple(new_key) From c7abefa4b7cc596848be28d22c633de4221f1485 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Mon, 23 Feb 2026 23:03:36 -0500 Subject: [PATCH 20/81] added tests for indexing --- src/finch/levels.py | 11 +++++ src/finch/scheduler.py | 2 +- src/finch/tensor.py | 44 +++++++++++++++--- tests/conftest.py | 35 ++++++++++++++ tests/test_indexing.py | 101 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 185 insertions(+), 8 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/test_indexing.py diff --git a/src/finch/levels.py b/src/finch/levels.py index 3a0ddfa..dde6d2a 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -25,6 +25,7 @@ class NestedLevelFType(LevelFType): def __init__(self, lvl: LevelFType): self.lvl = lvl + @property def ndim(self) -> np.intp: return self.lvl.ndim + np.intp(1) @@ -49,6 +50,7 @@ class Element(LevelFType): def __init__(self, fill_value: number): self._fill_value = fill_value + @property def ndim(self) -> np.intp: return np.intp(0) @@ -81,3 +83,12 @@ def create_jl_obj(self) -> JuliaObj: class SparseByteMap(NestedLevelFType): def create_jl_obj(self) -> JuliaObj: return jl.SparseByteMap(self.lvl.create_jl_obj()) + + +# Helper Methods +def construct_levels(obj: JuliaObj, fill_value: number) -> LevelFType: + if jl.isa(obj.lvl, jl.ElementLevel): + return Element(fill_value) + if jl.isa(obj.lvl, jl.DenseLevel): + return Dense(construct_levels(obj.lvl, fill_value)) + raise Exception("Unhandled exception!") diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py index 36b0289..fddfc31 100644 --- a/src/finch/scheduler.py +++ b/src/finch/scheduler.py @@ -11,7 +11,7 @@ from finchlite.finch_logic import LogicLoader from .compiler import FinchJLCompiler -from .tensor import Dense, Element +from .levels import Dense, Element class FinchJLLogicFormatter(LogicFormatter): diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 2a178f4..012b14f 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -4,8 +4,8 @@ from finchlite import EagerTensor, Tensor, TensorFType -from .julia import jl -from .levels import Dense, Element, NestedLevelFType +from .julia import jc, jl +from .levels import NestedLevelFType, construct_levels from .typing import JuliaObj from .utils import add_missing_dims, add_plus_one, expand_ellipsis @@ -15,6 +15,7 @@ class FinchJLTensorFType(TensorFType): def __init__(self, lvl): self._lvl: NestedLevelFType = lvl + @property def ndim(self) -> np.intp: return self._lvl.ndim @@ -49,14 +50,15 @@ def __init__(self, obj: JuliaObj): else: raise ValueError(f"Raw julia object expected. Found: {type(obj)}") - def ftype(self): + @property + def ftype(self) -> TensorFType: """Returns the ftype of the buffer""" - # TODO: figure out a way to walk through the levels and construct the ftype - return FinchJLTensorFType(Dense(Dense(Element(0)))) + return FinchJLTensorFType(construct_levels(self._obj, jl.fill_value(self._obj))) + @property def shape(self) -> tuple: """Shape of the tensor.""" - return self.obj.shape + return jl.size(self._obj) def __getitem__(self, key): if not isinstance(key, tuple): @@ -70,7 +72,35 @@ def __getitem__(self, key): result = self._obj[key] if jl.isa(result, jl.Finch.Tensor): return FinchJLTensor(result) - return result + return np.array(result) + + def _is_dense(self) -> bool: + lvl = self._obj.lvl + for _ in self.shape: + if not jl.isa(lvl, jl.Finch.Dense): + return False + lvl = lvl.lvl + return True + + def todense(self) -> np.ndarray: + obj = self._obj + + if self._is_dense: + # don't materialize a dense finch tensor + shape = jl.size(obj) + dense_tensor = obj.lvl + else: + # create materialized dense array + shape = jl.size(obj) + dense_lvls = jl.Element(jc.convert(self.dtype, jl.fill_value(obj))) + for _ in range(self.ndim): + dense_lvls = jl.Dense(dense_lvls) + dense_tensor = jl.Tensor(dense_lvls, obj).lvl # materialize + + for _ in range(self.ndim): + dense_tensor = dense_tensor.lvl + + return np.asarray(jl.reshape(dense_tensor.val, shape)) def __eq__(self, other): return isinstance(other, FinchJLTensor) and self._obj == other._obj diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..455e172 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,35 @@ +import pytest + +import numpy as np + + +@pytest.fixture +def rng(): + return np.random.default_rng(42) + + +@pytest.fixture +def arr1d(): + return np.arange(100) + + +@pytest.fixture +def arr2d(): + return np.array( + [ + [0, 0, 3, 2, 0], + [1, 0, 0, 1, 0], + [0, 5, 0, 0, 0], + ] + ) + + +@pytest.fixture +def arr3d(): + return np.array( + [ + [[0, 1, 0, 0], [1, 0, 0, 3]], + [[4, 0, -1, 0], [2, 2, 0, 0]], + [[0, 0, 0, 0], [1, 5, 0, 3]], + ] + ) diff --git a/tests/test_indexing.py b/tests/test_indexing.py new file mode 100644 index 0000000..5bd2c03 --- /dev/null +++ b/tests/test_indexing.py @@ -0,0 +1,101 @@ +import pytest + +from numpy.testing import assert_equal + +from juliacall import Main as jl + +from finch import FinchJLTensor + + +@pytest.mark.parametrize( + "index", + [ + 40, + (32,), + slice(30, 60, 3), + -10, + slice(None, -10, -2), + (None, slice(None)), + # The following two tests are commented out since Finch.jl + # returns errors for them + # + # ..., + # slice(None), + ], +) +def test_indexing_1d(arr1d, index): + arr_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Element(0)), arr1d)) + + actual = arr_finch[index] + expected = arr1d[index] + + if isinstance(actual, FinchJLTensor): + actual = actual.todense() + + assert_equal(actual, expected) + + +@pytest.mark.parametrize( + "index", + [ + ..., + 0, + (2,), + (2, 3), + slice(None), + (..., slice(0, 4, 2)), + (-1, slice(-1, None, -1)), + (None, slice(None), slice(None)), + ], +) +def test_indexing_2d(arr2d, index): + arr_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0))), arr2d)) + + actual = arr_finch[index] + expected = arr2d[index] + + if isinstance(actual, FinchJLTensor): + actual = actual.todense() + + assert_equal(actual, expected) + + +@pytest.mark.parametrize( + "index", + [ + (0, 1, 2), + (1, 0, 0), + (0, 1), + 1, + 2, + (2, slice(None), 3), + (slice(None), 0), + slice(None), + (0, slice(None), slice(1, 4, 2)), + (0, 1, ...), + (..., 1), + (0, ..., 1), + ..., + (..., slice(1, 4, 2)), + (slice(None, None, -1), slice(None, None, -1), slice(None, None, -1)), + (slice(None, -1, 1), slice(-1, None, -1), slice(4, 1, -1)), + (-1, 0, 0), + (0, -1, -2), + ([1, 2], 0, slice(3, None, -1)), + (0, slice(1, 0, -1), 0), + (slice(None), None, slice(None), slice(None)), + (slice(None), slice(None), slice(None), None), + ], +) +def test_indexing_3d(arr3d, index): + arr_finch = FinchJLTensor( + jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Dense(jl.Element(0)))), arr3d) + ) + + actual = arr_finch[index] + expected = arr3d[index] + + if isinstance(actual, FinchJLTensor): + actual = actual.todense() + + assert_equal(actual, expected) From 135424ea79f969e0093dd5f907cd61449fb2699f Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 24 Feb 2026 00:19:50 -0500 Subject: [PATCH 21/81] fix: resolved issues to start getting einsum tests to start working --- src/finch/compiler.py | 15 ++++++++++++--- src/finch/levels.py | 5 +++++ src/finch/scheduler.py | 3 ++- src/finch/tensor.py | 3 +++ src/finch/utils.py | 2 +- tests/test_einsum.py | 18 ++++++++++++++++++ 6 files changed, 41 insertions(+), 5 deletions(-) create mode 100644 tests/test_einsum.py diff --git a/src/finch/compiler.py b/src/finch/compiler.py index d9e3d9e..881719b 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,6 +1,7 @@ import operator import finchlite.finch_notation.nodes as ntn +from finchlite.algebra import make_tuple from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary @@ -8,6 +9,7 @@ from .tensor import FinchJLTensor ops_map = {operator.add: "+", operator.mul: "*"} +ops_to_ignore = [make_tuple] class FinchJLKernel(AssemblyKernel): @@ -69,7 +71,9 @@ def generate_julia(self, prgm, nestingLvl=0): # TODO: Can we make this better? # Special condition to ignore all assigns associated with # finding loop bounds - if isinstance(rhs, ntn.Call) and rhs.op.val == dimension: + if isinstance(rhs, ntn.Dimension) or ( + isinstance(rhs, ntn.Call) and rhs.op.val == dimension + ): return "" tab_str = " " * nestingLvl @@ -121,6 +125,8 @@ def generate_julia(self, prgm, nestingLvl=0): arg_str = ",".join( [self.generate_julia(arg, nestingLvl) for arg in args] ) + if op.val in ops_to_ignore: + return f"{arg_str}" return f"{ops_map[op.val]}({arg_str})" case ntn.If(cond, body): @@ -177,7 +183,9 @@ def generate_julia(self, prgm, nestingLvl=0): return str(val) case ntn.Variable(name, _): - return name + # finch tensor lite uses character(#) in the naming of variables + # that however is not valid julia syntax + return name.replace("#", "_") # TODO: Cached, Dimension, Thaw, Stack, Value are unimplemented. case _: @@ -190,6 +198,7 @@ def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: kernel_dict = {} for func in prgm.children: - kernel_dict[func.name.name] = FinchJLKernel(func.name.name, generator(func)) + generated_prgm = generator(func) + kernel_dict[func.name.name] = FinchJLKernel(func.name.name, generated_prgm) return FinchJLLibrary(kernel_dict) diff --git a/src/finch/levels.py b/src/finch/levels.py index dde6d2a..92a13c8 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -14,6 +14,7 @@ class LevelFType(TensorFType): def from_numpy(self, _) -> Tensor: raise NotImplementedError + @property def shape_type(self) -> tuple[type, ...]: return tuple(self.element_type for _ in range(self.ndim)) @@ -29,9 +30,11 @@ def __init__(self, lvl: LevelFType): def ndim(self) -> np.intp: return self.lvl.ndim + np.intp(1) + @property def fill_value(self) -> Any: return self.lvl.fill_value + @property def element_type(self) -> Any: return self.lvl.element_type @@ -54,9 +57,11 @@ def __init__(self, fill_value: number): def ndim(self) -> np.intp: return np.intp(0) + @property def fill_value(self) -> Any: return self._fill_value + @property def element_type(self) -> Any: return type(self._fill_value) diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py index fddfc31..7a0cbb6 100644 --- a/src/finch/scheduler.py +++ b/src/finch/scheduler.py @@ -12,6 +12,7 @@ from .compiler import FinchJLCompiler from .levels import Dense, Element +from .tensor import FinchJLTensorFType class FinchJLLogicFormatter(LogicFormatter): @@ -25,7 +26,7 @@ def get_output_tns_ftype(self, fill_value: Any, shape_type: tuple[Any, ...]): lvl = Element(fill_value) for _ in len(shape_type): lvl = Dense(lvl) - return lvl + return FinchJLTensorFType(lvl) COMPILE_JULIA = LogicNormalizer( diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 012b14f..ba4ccfe 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -19,12 +19,15 @@ def __init__(self, lvl): def ndim(self) -> np.intp: return self._lvl.ndim + @property def fill_value(self) -> Any: return self._lvl.fill_value + @property def element_type(self) -> Any: return self._lvl.element_type + @property def shape_type(self) -> tuple[type, ...]: return self._lvl.shape_type diff --git a/src/finch/utils.py b/src/finch/utils.py index accfa47..acdf3fb 100644 --- a/src/finch/utils.py +++ b/src/finch/utils.py @@ -3,7 +3,7 @@ import builtins import numpy as np -from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple +from numpy._core.numeric import normalize_axis_index, normalize_axis_tuple from .julia import jl diff --git a/tests/test_einsum.py b/tests/test_einsum.py new file mode 100644 index 0000000..61e6a6b --- /dev/null +++ b/tests/test_einsum.py @@ -0,0 +1,18 @@ +import numpy as np + +import finchlite +from juliacall import Main as jl + +from finch import COMPILE_JULIA, FinchJLTensor + + +def test_pass_through(rng): + """Test pass through of a tensor""" + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A = FinchJLTensor( + jl.Finch.Tensor( + jl.Dense(jl.Dense(jl.Element(0.0))), np.array(rng.random((5, 5))) + ) + ) + B = finchlite.einop("B[i,j] = A[i,j]", A=A) + np.allclose(B.todense(), A.todense()) From 60ccc6474a198ee991873c8d95047405940290dc Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 24 Feb 2026 10:11:54 -0500 Subject: [PATCH 22/81] feat: passed transpose einsum test --- src/finch/compiler.py | 13 +++++++++---- src/finch/scheduler.py | 2 +- tests/test_einsum.py | 24 +++++++++++++++++------- 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 881719b..d3b5cfc 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -2,13 +2,14 @@ import finchlite.finch_notation.nodes as ntn from finchlite.algebra import make_tuple +from finchlite.algebra.operator import overwrite from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary from .julia import jl from .tensor import FinchJLTensor -ops_map = {operator.add: "+", operator.mul: "*"} +ops_map = {operator.add: "+", operator.mul: "*", operator.eq: "=="} ops_to_ignore = [make_tuple] @@ -97,6 +98,7 @@ def generate_julia(self, prgm, nestingLvl=0): case ntn.Loop(idx, _, body): tab_str = " " * nestingLvl tab_str_1 = " " * (nestingLvl + 1) + idx_str = self.generate_julia(idx, nestingLvl) is_outermost_loop = False if self.in_finch_block is False: @@ -107,10 +109,10 @@ def generate_julia(self, prgm, nestingLvl=0): loop_body = self.generate_julia(body, nestingLvl + 1) if not is_outermost_loop: - return f"{tab_str}for {idx.name} = _\n{loop_body}{tab_str}end\n" + return f"{tab_str}for {idx_str} = _\n{loop_body}{tab_str}end\n" self.in_finch_block = False return ( - f"{tab_str}@finch begin\n{tab_str_1}for {idx.name} = " + f"{tab_str}@finch begin\n{tab_str_1}for {idx_str} = " f"_\n{loop_body}{tab_str_1}end\n{tab_str}end" ) @@ -155,6 +157,9 @@ def generate_julia(self, prgm, nestingLvl=0): ): raise Exception("Increment expects the lhs to be an access") + # If the operation is overwrite just codegen an assignment + if lhs.mode.op.val == overwrite: + return f"{tab_str}{lhs_str} = {rhs_str}" return f"{tab_str}{lhs_str} {ops_map[lhs.mode.op.val]}= {rhs_str}" case ntn.Unwrap(arg): @@ -164,7 +169,7 @@ def generate_julia(self, prgm, nestingLvl=0): # TODO: Is this the right assumption to make if not isinstance(rhs, ntn.Variable): raise Exception("The unpack was not called with variable as RHS.") - self.pack_dict[lhs.name] = rhs.name + self.pack_dict[lhs.name] = self.generate_julia(rhs, nestingLvl) return "" case ntn.Repack(val, _): diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py index 7a0cbb6..d9f4d01 100644 --- a/src/finch/scheduler.py +++ b/src/finch/scheduler.py @@ -24,7 +24,7 @@ def __init__( def get_output_tns_ftype(self, fill_value: Any, shape_type: tuple[Any, ...]): lvl = Element(fill_value) - for _ in len(shape_type): + for _ in shape_type: lvl = Dense(lvl) return FinchJLTensorFType(lvl) diff --git a/tests/test_einsum.py b/tests/test_einsum.py index 61e6a6b..bc57352 100644 --- a/tests/test_einsum.py +++ b/tests/test_einsum.py @@ -8,11 +8,21 @@ def test_pass_through(rng): """Test pass through of a tensor""" + A = rng.random((5, 5)) + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A = FinchJLTensor( - jl.Finch.Tensor( - jl.Dense(jl.Dense(jl.Element(0.0))), np.array(rng.random((5, 5))) - ) - ) - B = finchlite.einop("B[i,j] = A[i,j]", A=A) - np.allclose(B.todense(), A.todense()) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + B = finchlite.einop("B[i,j] = A[i,j]", A=A_finch) + + np.allclose(B.todense(), A) + + +def test_transpose(rng): + """Test basic addition with transpose""" + A = rng.random((5, 5)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + B = finchlite.einop("B[i,j] = A[j, i]", A=A_finch) + + np.allclose(B.todense(), A.T) From fd68a1fcb14f4243cddebbea1f5f63918f34b259 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 24 Feb 2026 10:54:31 -0500 Subject: [PATCH 23/81] feat: expanded covered einsum tests --- src/finch/compiler.py | 16 +++++- tests/test_einsum.py | 110 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 124 insertions(+), 2 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index d3b5cfc..aef3c7a 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,8 +1,9 @@ +import math import operator import finchlite.finch_notation.nodes as ntn from finchlite.algebra import make_tuple -from finchlite.algebra.operator import overwrite +from finchlite.algebra.operator import overwrite, promote_max, promote_min from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary @@ -10,6 +11,12 @@ from .tensor import FinchJLTensor ops_map = {operator.add: "+", operator.mul: "*", operator.eq: "=="} +red_ops_map = { + operator.add: "+", + operator.mul: "*", + promote_max: "<>", + promote_min: "<>", +} ops_to_ignore = [make_tuple] @@ -160,7 +167,7 @@ def generate_julia(self, prgm, nestingLvl=0): # If the operation is overwrite just codegen an assignment if lhs.mode.op.val == overwrite: return f"{tab_str}{lhs_str} = {rhs_str}" - return f"{tab_str}{lhs_str} {ops_map[lhs.mode.op.val]}= {rhs_str}" + return f"{tab_str}{lhs_str} {red_ops_map[lhs.mode.op.val]}= {rhs_str}" case ntn.Unwrap(arg): return self.generate_julia(arg, nestingLvl) @@ -185,6 +192,11 @@ def generate_julia(self, prgm, nestingLvl=0): return self.pack_dict[name] case ntn.Literal(val): + # Julia represents inf differently than how its represented in python + if val > 0 and math.isinf(val): + return "Inf" + if val < 0 and math.isinf(val): + return "-Inf" return str(val) case ntn.Variable(name, _): diff --git a/tests/test_einsum.py b/tests/test_einsum.py index bc57352..32f92f0 100644 --- a/tests/test_einsum.py +++ b/tests/test_einsum.py @@ -26,3 +26,113 @@ def test_transpose(rng): B = finchlite.einop("B[i,j] = A[j, i]", A=A_finch) np.allclose(B.todense(), A.T) + + +def test_basic_addition_with_transpose(rng): + """Test basic addition with transpose""" + A = rng.random((5, 5)) + B = rng.random((5, 5)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), B)) + C = finchlite.einop("C[i,j] = A[i,j] + B[j,i]", A=A_finch, B=B_finch) + C_ref = A + B.T + + np.allclose(C.todense(), C_ref) + + +def test_matrix_multiplication(rng): + """Test matrix multiplication using += (increment/accumulation)""" + A = rng.random((3, 4)) + B = rng.random((4, 5)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), B)) + C = finchlite.einop("C[i,j] += A[i,k] * B[k,j]", A=A_finch, B=B_finch) + C_ref = A @ B + + np.allclose(C.todense(), C_ref) + + +def test_element_wise_multiplication(rng): + """Test element-wise multiplication""" + A = rng.random((4, 4)) + B = rng.random((4, 4)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), B)) + C = finchlite.einop("C[i,j] = A[i,j] * B[i,j]", A=A_finch, B=B_finch) + C_ref = A * B + + np.allclose(C.todense(), C_ref) + + +def test_sum_reduction(rng): + """Test sum reduction using +=""" + A = rng.random((3, 4)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + C = finchlite.einop("C[i] += A[i,j]", A=A_finch) + C_ref = np.sum(A, axis=1) + + np.allclose(C.todense(), C_ref) + + +def test_maximum_reduction(rng): + """Test maximum reduction using max=""" + A = rng.random((3, 4)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + C = finchlite.einop("C[i] max= A[i,j]", A=A_finch) + C_ref = np.max(A, axis=1) + + np.allclose(C.todense(), C_ref) + + +def test_outer_product(rng): + """Test outer product""" + A = rng.random(3) + B = rng.random(4) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Element(0.0)), A)) + B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Element(0.0)), B)) + C = finchlite.einop("C[i,j] = A[i] * B[j]", A=A_finch, B=B_finch) + C_ref = np.outer(A, B) + + np.allclose(C.todense(), C_ref) + + +def test_batch_matrix_multiplication(rng): + """Test batch matrix multiplication using +=""" + A = rng.random((2, 3, 4)) + B = rng.random((2, 4, 5)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor( + jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Dense(jl.Element(0.0)))), A) + ) + B_finch = FinchJLTensor( + jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Dense(jl.Element(0.0)))), B) + ) + C = finchlite.einop("C[b,i,j] += A[b,i,k] * B[b,k,j]", A=A_finch, B=B_finch) + C_ref = np.matmul(A, B) + + np.allclose(C.todense(), C_ref) + + +def test_minimum_reduction(rng): + """Test minimum reduction using min=""" + A = rng.random((3, 4)) + + finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) + A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + C = finchlite.einop("C[i] min= A[i,j]", A=A_finch) + C_ref = np.min(A, axis=1) + + np.allclose(C.todense(), C_ref) From 362cd2f1b9ab2274bedfdfae44e60322678fd66a Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 24 Feb 2026 11:52:04 -0500 Subject: [PATCH 24/81] feat: match cases for finch.jl level matching --- src/finch/levels.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 92a13c8..5533386 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -92,8 +92,12 @@ def create_jl_obj(self) -> JuliaObj: # Helper Methods def construct_levels(obj: JuliaObj, fill_value: number) -> LevelFType: - if jl.isa(obj.lvl, jl.ElementLevel): + if jl.isa(obj.lvl, jl.Finch.Element): return Element(fill_value) - if jl.isa(obj.lvl, jl.DenseLevel): + if jl.isa(obj.lvl, jl.Finch.Dense): return Dense(construct_levels(obj.lvl, fill_value)) + if jl.isa(obj.lvl, jl.Finch.SparseList): + return SparseList(construct_levels(obj.lvl, fill_value)) + if jl.isa(obj.lvl, jl.Finch.SparseByteMap): + return SparseByteMap(construct_levels(obj.lvl, fill_value)) raise Exception("Unhandled exception!") From 74019c4baf4d6cef589925f6aaced5ab84b4847d Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 25 Feb 2026 13:51:30 -0500 Subject: [PATCH 25/81] feat: added finch-tensor-lite as a dependency --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 59a86d3..6925204 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "numpy (>=1.19,<2.4)", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", + "finch-tensor-lite (==0.3.0)", ] [tool.poetry] From 45473a82f2da89d797e82950f20453b8b6eac5a7 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 25 Feb 2026 14:11:58 -0500 Subject: [PATCH 26/81] fix: import bug --- src/finch/__init__.py | 5 ----- tests/test_compiler.py | 4 ++-- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/src/finch/__init__.py b/src/finch/__init__.py index 6ef8cdd..c16f141 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -1,4 +1,3 @@ -from .levels import Dense, Element, SparseByteMap, SparseList from .scheduler import COMPILE_JULIA from .tensor import ( FinchJLTensor, @@ -7,10 +6,6 @@ __all__ = [ "COMPILE_JULIA", - "Dense", - "Element", "FinchJLTensor", "FinchJLTensorFType", - "SparseByteMap", - "SparseList", ] diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 3517ba8..0191f69 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -29,12 +29,12 @@ from finch.compiler import FinchJLCompiler, FinchJLKernel from finch.julia import jl -from finch.tensor import Dense, Element, FinchJLTensor +from finch.levels import Dense, Element +from finch.tensor import FinchJLTensor a_format = Dense(Dense(Element(0))) -@pytest.mark.skip @pytest.mark.parametrize( "finch_ntn, julia_code", [ From b8566ebc3a946f58829351c7b4accd307313a7fd Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 25 Feb 2026 15:18:58 -0500 Subject: [PATCH 27/81] feat: adding back array api tests --- src/finch/__init__.py | 246 +++++++++++++++++++++++++++++++++++ src/finch/_array_api_info.py | 94 +++++++++++++ src/finch/dtypes.py | 61 +++++++++ 3 files changed, 401 insertions(+) create mode 100644 src/finch/_array_api_info.py create mode 100644 src/finch/dtypes.py diff --git a/src/finch/__init__.py b/src/finch/__init__.py index c16f141..40b13e9 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -1,3 +1,125 @@ +from finchlite import ( + abs, + acos, + acosh, + add, + all, + any, + asin, + asinh, + atan, + atan2, + atanh, + bitwise_and, + bitwise_inverse, + bitwise_left_shift, + bitwise_or, + bitwise_right_shift, + bitwise_xor, + broadcast_arrays, + broadcast_to, + ceil, + clip, + combine_dims, + concat, + copysign, + cos, + cosh, + divide, + einop, + einsum, + elementwise, + equal, + exp, + expand_dims, + expm1, + flatten, + floor, + floordiv, + greater, + greater_equal, + hypot, + imag, + isfinite, + isinf, + isnan, + less, + less_equal, + log, + log1p, + log2, + log10, + logaddexp, + logical_and, + logical_not, + logical_or, + logical_xor, + matmul, + matrix_transpose, + max, + maximum, + mean, + min, + minimum, + mod, + moveaxis, + multiply, + negative, + nextafter, + not_equal, + permute_dims, + positive, + pow, + power, + prod, + real, + reciprocal, + reduce, + remainder, + round, + sign, + signbit, + sin, + sinh, + split_dims, + sqrt, + square, + squeeze, + stack, + std, + subtract, + sum, + tan, + tanh, + tensordot, + truediv, + trunc, + var, + vecdot, +) + +from ._array_api_info import __array_namespace_info__ +from .dtypes import ( + bool, + can_cast, + complex64, + complex128, + finfo, + float16, + float32, + float64, + iinfo, + int8, + int16, + int32, + int64, + int_, + uint, + uint8, + uint16, + uint32, + uint64, +) from .scheduler import COMPILE_JULIA from .tensor import ( FinchJLTensor, @@ -8,4 +130,128 @@ "COMPILE_JULIA", "FinchJLTensor", "FinchJLTensorFType", + "__array_namespace_info__", + "abs", + "acos", + "acosh", + "add", + "all", + "any", + "asarray", + "asin", + "asinh", + "atan", + "atan2", + "atanh", + "bitwise_and", + "bitwise_inverse", + "bitwise_left_shift", + "bitwise_or", + "bitwise_right_shift", + "bitwise_xor", + "bool", + "broadcast_arrays", + "broadcast_to", + "can_cast", + "ceil", + "clip", + "combine_dims", + "complex64", + "complex128", + "compute", + "concat", + "copysign", + "cos", + "cosh", + "divide", + "einop", + "einsum", + "elementwise", + "equal", + "exp", + "expand_dims", + "expm1", + "finfo", + "flatten", + "float16", + "float32", + "float64", + "floor", + "floordiv", + "fuse", + "fused", + "get_default_scheduler", + "greater", + "greater_equal", + "hypot", + "iinfo", + "imag", + "int8", + "int16", + "int32", + "int64", + "int_", + "isfinite", + "isinf", + "isnan", + "lazy", + "less", + "less_equal", + "log", + "log1p", + "log2", + "log10", + "logaddexp", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "matmul", + "matrix_transpose", + "max", + "maximum", + "mean", + "min", + "minimum", + "mod", + "moveaxis", + "multiply", + "negative", + "nextafter", + "not_equal", + "permute_dims", + "positive", + "pow", + "power", + "prod", + "real", + "reciprocal", + "reduce", + "remainder", + "round", + "set_default_scheduler", + "sign", + "signbit", + "sin", + "sinh", + "split_dims", + "sqrt", + "square", + "squeeze", + "stack", + "std", + "subtract", + "sum", + "tan", + "tanh", + "tensordot", + "truediv", + "trunc", + "uint", + "uint8", + "uint16", + "uint32", + "uint64", + "var", + "vecdot", ] diff --git a/src/finch/_array_api_info.py b/src/finch/_array_api_info.py new file mode 100644 index 0000000..c4e3f5a --- /dev/null +++ b/src/finch/_array_api_info.py @@ -0,0 +1,94 @@ +from . import dtypes +from .typing import DType + + +class __array_namespace_info__: + def capabilities(self) -> dict[str, bool]: + return { + "boolean indexing": True, + "data-dependent shapes": True, + } + + def default_device(self) -> str: + return "cpu" + + def default_dtypes(self, *, device: str | None = None) -> dict[str, DType]: + if device not in ["cpu", None]: + raise ValueError( + f'Device not understood. Only "cpu" is allowed, but received: {device}' + ) + return { + "real floating": dtypes.float64, + "complex floating": dtypes.complex128, + "integral": dtypes.int_, + "indexing": dtypes.int_, + } + + _bool_dtypes = {"bool": dtypes.bool} + _signed_integer_dtypes = { + "int8": dtypes.int8, + "int16": dtypes.int16, + "int32": dtypes.int32, + "int64": dtypes.int64, + } + _unsigned_integer_dtypes = { + "uint8": dtypes.uint8, + "uint16": dtypes.uint16, + "uint32": dtypes.uint32, + "uint64": dtypes.uint64, + } + _real_floating_dtypes = { + "float32": dtypes.float32, + "float64": dtypes.float64, + } + _complex_floating_dtypes = { + "complex64": dtypes.complex64, + "complex128": dtypes.complex128, + } + + def dtypes( + self, + *, + device: str | None = None, + kind: str | tuple[str, ...] | None = None, + ) -> dict[str, DType]: + if device not in ["cpu", None]: + raise ValueError( + f'Device not understood. Only "cpu" is allowed, but received: {device}' + ) + if kind is None: + return ( + self._bool_dtypes + | self._signed_integer_dtypes + | self._unsigned_integer_dtypes + | self._real_floating_dtypes + | self._complex_floating_dtypes + ) + if kind == "bool": + return self._bool_dtypes + if kind == "signed integer": + return self._signed_integer_dtypes + if kind == "unsigned integer": + return self._unsigned_integer_dtypes + if kind == "integral": + return self._signed_integer_dtypes | self._unsigned_integer_dtypes + if kind == "real floating": + return self._real_floating_dtypes + if kind == "complex floating": + return self._complex_floating_dtypes + if kind == "numeric": + return ( + self._signed_integer_dtypes + | self._unsigned_integer_dtypes + | self._real_floating_dtypes + | self._complex_floating_dtypes + ) + if isinstance(kind, tuple): + res = {} + for k in kind: + res.update(self.dtypes(kind=k)) + return res + raise ValueError(f"unsupported kind: {kind!r}") + + def devices(self) -> list[str]: + return ["cpu"] diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py new file mode 100644 index 0000000..2b95252 --- /dev/null +++ b/src/finch/dtypes.py @@ -0,0 +1,61 @@ +import builtins + +import numpy as np + +from .julia import jl + +int_: jl.DataType = jl.Int +int8: jl.DataType = jl.Int8 +int16: jl.DataType = jl.Int16 +int32: jl.DataType = jl.Int32 +int64: jl.DataType = jl.Int64 +uint: jl.DataType = jl.UInt +uint8: jl.DataType = jl.UInt8 +uint16: jl.DataType = jl.UInt16 +uint32: jl.DataType = jl.UInt32 +uint64: jl.DataType = jl.UInt64 +float16: jl.DataType = jl.Float16 +float32: jl.DataType = jl.Float32 +float64: jl.DataType = jl.Float64 +complex64: jl.DataType = jl.ComplexF32 +complex128: jl.DataType = jl.ComplexF64 +bool: jl.DataType = jl.Bool + +number: jl.DataType = jl.Number +complex: jl.DataType = jl.Complex +integer: jl.DataType = jl.Integer +abstract_float: jl.DataType = jl.AbstractFloat + +jl_to_np_dtype = { + int_: np.int_, + int8: np.int8, + int16: np.int16, + int32: np.int32, + int64: np.int64, + uint: np.uint, + uint8: np.uint8, + uint16: np.uint16, + uint32: np.uint32, + uint64: np.uint64, + float16: np.float16, + float32: np.float32, + float64: np.float64, + complex64: np.complex64, + complex128: np.complex128, + bool: builtins.bool, + None: None, +} + + +def finfo(dtype): + return np.finfo(jl_to_np_dtype[dtype]) + + +def iinfo(dtype): + return np.iinfo(jl_to_np_dtype[dtype]) + + +def can_cast(from_, to, /) -> builtins.bool: + if hasattr(from_, "dtype"): + from_ = from_.dtype + return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) From 61aa6fcac1adca56b98cdc646d44a8600514c12b Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 25 Feb 2026 17:13:15 -0500 Subject: [PATCH 28/81] fix: reintroducing dtype --- src/finch/typing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/finch/typing.py b/src/finch/typing.py index a5bacb6..d3003ab 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -1,5 +1,5 @@ import juliacall as jc JuliaObj = jc.AnyValue - +DType = jc.AnyValue number = int | float | bool | complex From 2ab119bc935909698cdbcbab4caf169b57b6fc70 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 25 Feb 2026 17:43:41 -0500 Subject: [PATCH 29/81] chore: updating gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 9167885..87ddf79 100644 --- a/.gitignore +++ b/.gitignore @@ -51,6 +51,7 @@ coverage.xml .pytest_cache/ cover/ junit/ +array-api-tests/ # Translations *.mo From 94f03acccd9d35e412a0544770daba4fa8ec3b45 Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Wed, 11 Mar 2026 11:57:25 -0400 Subject: [PATCH 30/81] feat: adding scalars --- src/finch/levels.py | 49 +++++++++++++++++++++++++++++++++++---------- src/finch/tensor.py | 25 ++++++++++++++++++++--- 2 files changed, 60 insertions(+), 14 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 5533386..5f85849 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -22,33 +22,32 @@ def __call__(self, _) -> Tensor: raise Exception("Cannot create an object of this type!") -class NestedLevelFType(LevelFType): - def __init__(self, lvl: LevelFType): - self.lvl = lvl +class Scalar(LevelFType): + def __init__(self, val: number): + self._val = val @property def ndim(self) -> np.intp: - return self.lvl.ndim + np.intp(1) + return np.intp(0) @property def fill_value(self) -> Any: - return self.lvl.fill_value + return self._val @property def element_type(self) -> Any: - return self.lvl.element_type + return type(self._val) def __eq__(self, other): - return type(other) is type(self) and self.lvl == other.lvl + return isinstance(other, Scalar) and self._val == other._val def __hash__(self): - return hash((self.__class__.__name__, self.lvl.__hash__)) + return hash((self.__class__.__name__, self._val)) - @abstractmethod - def create_jl_obj(self) -> JuliaObj: ... + def create_jl_obj(self) -> JuliaObj: + return jl.Scalar(self._val) -# Concrete FTypes class Element(LevelFType): def __init__(self, fill_value: number): self._fill_value = fill_value @@ -75,6 +74,34 @@ def create_jl_obj(self) -> JuliaObj: return jl.Element(self._fill_value) +class NestedLevelFType(LevelFType): + def __init__(self, lvl: LevelFType): + if not isinstance(lvl, NestedLevelFType | Element): + raise ValueError("lvl must be a NestedLevelFType or Element.") + self.lvl = lvl + + @property + def ndim(self) -> np.intp: + return self.lvl.ndim + np.intp(1) + + @property + def fill_value(self) -> Any: + return self.lvl.fill_value + + @property + def element_type(self) -> Any: + return self.lvl.element_type + + def __eq__(self, other): + return type(other) is type(self) and self.lvl == other.lvl + + def __hash__(self): + return hash((self.__class__.__name__, self.lvl.__hash__)) + + @abstractmethod + def create_jl_obj(self) -> JuliaObj: ... + + class Dense(NestedLevelFType): def create_jl_obj(self) -> JuliaObj: return jl.Dense(self.lvl.create_jl_obj()) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index ba4ccfe..01311da 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -5,7 +5,7 @@ from finchlite import EagerTensor, Tensor, TensorFType from .julia import jc, jl -from .levels import NestedLevelFType, construct_levels +from .levels import LevelFType, Scalar, construct_levels from .typing import JuliaObj from .utils import add_missing_dims, add_plus_one, expand_ellipsis @@ -13,7 +13,7 @@ # Tensor Class and associated ftype class FinchJLTensorFType(TensorFType): def __init__(self, lvl): - self._lvl: NestedLevelFType = lvl + self._lvl: LevelFType = lvl @property def ndim(self) -> np.intp: @@ -31,7 +31,12 @@ def element_type(self) -> Any: def shape_type(self) -> tuple[type, ...]: return self._lvl.shape_type - def __call__(self, shape: tuple) -> Tensor: + def __call__(self, shape: tuple | None = None) -> Tensor: + if isinstance(self._lvl, Scalar): + return FinchJLTensor(self._lvl.create_jl_obj()) + + if shape is None: + raise ValueError("shape argument cannot be None for non scalar tensors.") return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), shape)) def from_numpy(self, _) -> Tensor: @@ -56,6 +61,8 @@ def __init__(self, obj: JuliaObj): @property def ftype(self) -> TensorFType: """Returns the ftype of the buffer""" + if self._is_scalar(): + return FinchJLTensorFType(Scalar(self._obj.val)) return FinchJLTensorFType(construct_levels(self._obj, jl.fill_value(self._obj))) @property @@ -64,6 +71,9 @@ def shape(self) -> tuple: return jl.size(self._obj) def __getitem__(self, key): + if self._is_scalar(): + raise ValueError("Scalars are not subscriptable!") + if not isinstance(key, tuple): key = (key,) @@ -77,7 +87,13 @@ def __getitem__(self, key): return FinchJLTensor(result) return np.array(result) + def _is_scalar(self) -> bool: + return jl.isa(self._obj, jl.Finch.Scalar) + def _is_dense(self) -> bool: + if self._is_scalar(): + return False + lvl = self._obj.lvl for _ in self.shape: if not jl.isa(lvl, jl.Finch.Dense): @@ -86,6 +102,9 @@ def _is_dense(self) -> bool: return True def todense(self) -> np.ndarray: + if self._is_scalar(): + return np.asarray(self._obj.val) + obj = self._obj if self._is_dense: From 74a4b2397a561aa1202ddbce35e7efc65c673ecf Mon Sep 17 00:00:00 2001 From: JoelMathewC Date: Tue, 17 Mar 2026 11:41:30 -0400 Subject: [PATCH 31/81] fix: broken import --- pyproject.toml | 2 +- src/finch/compiler.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6925204..40398ef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "numpy (>=1.19,<2.4)", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", - "finch-tensor-lite (==0.3.0)", + # "finch-tensor-lite (==0.3.0)", ] [tool.poetry] diff --git a/src/finch/compiler.py b/src/finch/compiler.py index aef3c7a..433e7ee 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -3,7 +3,7 @@ import finchlite.finch_notation.nodes as ntn from finchlite.algebra import make_tuple -from finchlite.algebra.operator import overwrite, promote_max, promote_min +from finchlite.algebra.algebra import overwrite, promote_max, promote_min from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary From 5ad4fa302d35eb96fe49f50da8015bd9cdbd0129 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 14:03:10 -0400 Subject: [PATCH 32/81] fixed --- array-api-tests | 1 + pyproject.toml | 3 ++ {ci => tests}/array-api-skips.txt | 0 tests/test_array_api.py | 63 +++++++++++++++++++++++++++++++ 4 files changed, 67 insertions(+) create mode 160000 array-api-tests rename {ci => tests}/array-api-skips.txt (100%) create mode 100644 tests/test_array_api.py diff --git a/array-api-tests b/array-api-tests new file mode 160000 index 0000000..c48410f --- /dev/null +++ b/array-api-tests @@ -0,0 +1 @@ +Subproject commit c48410f96fc58e02eea844e6b7f6cc01680f77ce diff --git a/pyproject.toml b/pyproject.toml index 999b30c..b6a1e02 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,9 @@ test = [ "pytest-cov>=4.1.0,<4.2", "sparse>=0.17.0,<0.17.1", "scipy>=1.7,<1.17", + "hypothesis (>=6.151.9,<7.0.0)", + "ndindex (>=1.10.1,<2.0.0)", + "pytest-json-report (>=1.5.0,<2.0.0)", ] [build-system] diff --git a/ci/array-api-skips.txt b/tests/array-api-skips.txt similarity index 100% rename from ci/array-api-skips.txt rename to tests/array-api-skips.txt diff --git a/tests/test_array_api.py b/tests/test_array_api.py new file mode 100644 index 0000000..77311f2 --- /dev/null +++ b/tests/test_array_api.py @@ -0,0 +1,63 @@ +import os +import subprocess +import sys + + +def test_array_api(): + ARRAY_API_TESTS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), os.environ.get("ARRAY_API_TESTS_DIR", "../array-api-tests"))) + ARRAY_API_TESTS_REV = os.environ.get("ARRAY_API_TESTS_REV", "c48410f96fc58e02eea844e6b7f6cc01680f77ce") + ARRAY_API_TESTS_SKIPS = os.path.abspath(os.path.join(os.path.dirname(__file__), os.environ.get("ARRAY_API_TESTS_SKIPS", "array-api-skips.txt"))) + ARRAY_API_TESTS_ARGS = os.environ.get("ARRAY_API_TESTS_ARGS", "-vv -s") + + print(f"[array-api] using dir: {ARRAY_API_TESTS_DIR}", flush=True) + print(f"[array-api] target rev: {ARRAY_API_TESTS_REV}", flush=True) + + if not os.path.isdir(ARRAY_API_TESTS_DIR): + print("[array-api] cloning test repo...", flush=True) + subprocess.run( + ["git", "clone", "--recursive", + "https://github.com/data-apis/array-api-tests.git", + ARRAY_API_TESTS_DIR], + check=True + ) + + print("[array-api] cleaning repo...", flush=True) + subprocess.run( + ["git", "--git-dir", f"{ARRAY_API_TESTS_DIR}/.git", + "--work-tree", ARRAY_API_TESTS_DIR, "clean", "-xddf"], + check=True + ) + + print("[array-api] fetching latest refs...", flush=True) + subprocess.run( + ["git", "--git-dir", f"{ARRAY_API_TESTS_DIR}/.git", + "--work-tree", ARRAY_API_TESTS_DIR, "fetch"], + check=True + ) + + print("[array-api] checking out target revision...", flush=True) + subprocess.run( + ["git", "--git-dir", f"{ARRAY_API_TESTS_DIR}/.git", + "--work-tree", ARRAY_API_TESTS_DIR, "reset", "--hard", ARRAY_API_TESTS_REV], + check=True + ) + + # Run the tests using pytest + print("[array-api] running external array-api-tests...", flush=True) + result = subprocess.run( + [ + sys.executable, "-m", "pytest", + *ARRAY_API_TESTS_ARGS.split(), + f"{ARRAY_API_TESTS_DIR}/array_api_tests/", + "--max-examples=2", + "--derandomize", + "--disable-deadline", + "--disable-warnings", + "--skips-file", ARRAY_API_TESTS_SKIPS + ], + env={**os.environ, "ARRAY_API_TESTS_MODULE": "finch", "PYTHONUNBUFFERED": "1"}, + check=False, + text=True, + ) + print("[array-api] array-api-tests completed!", flush=True) + assert result.returncode == 0, "Array API tests failed" \ No newline at end of file From f9750648011b4f63b9a169f0b94243781ede68e0 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 14:11:29 -0400 Subject: [PATCH 33/81] cool --- .github/workflows/ci.yml | 32 -------------------------------- .gitignore | 1 + array-api-tests | 1 - ci/array-api-tests-rev.txt | 1 - ci/clone_array_api_tests.sh | 11 ----------- 5 files changed, 1 insertion(+), 45 deletions(-) delete mode 160000 array-api-tests delete mode 100644 ci/array-api-tests-rev.txt delete mode 100755 ci/clone_array_api_tests.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 70e150e..7f1e00f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -38,38 +38,6 @@ jobs: poetry run pytest --junit-xml=test-${{ matrix.os }}-Python-${{ matrix.python }}.xml - uses: codecov/codecov-action@v3 - array_api_tests: - env: - ARRAY_API_TESTS_DIR: ${{ github.workspace }}/array-api-tests - runs-on: ubuntu-latest - steps: - - name: Checkout Repo - uses: actions/checkout@v4 - - name: Checkout array-api-tests - run: ci/clone_array_api_tests.sh - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.11' - cache: 'pip' - - name: Install Poetry - run: | - pip install poetry - - name: Build wheel - run: | - python -m poetry build --format wheel - - name: Install build and test dependencies from PyPI - run: | - pip install dist/*.whl - pip install -U setuptools wheel - pip install pytest-xdist hypothesis==6.131.0 -r "$ARRAY_API_TESTS_DIR/requirements.txt" - - name: Run the test suite - env: - ARRAY_API_TESTS_MODULE: finch - run: | - python -c 'import finch' - pytest "$ARRAY_API_TESTS_DIR/array_api_tests/" -v -c "$ARRAY_API_TESTS_DIR/pytest.ini" --ci --max-examples=2 --derandomize --disable-deadline --disable-warnings -n auto --skips-file ci/array-api-skips.txt - on: # Trigger the workflow on push or pull request, # but only for the main branch diff --git a/.gitignore b/.gitignore index 4896382..2b2b89f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +array-api-tests # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/array-api-tests b/array-api-tests deleted file mode 160000 index c48410f..0000000 --- a/array-api-tests +++ /dev/null @@ -1 +0,0 @@ -Subproject commit c48410f96fc58e02eea844e6b7f6cc01680f77ce diff --git a/ci/array-api-tests-rev.txt b/ci/array-api-tests-rev.txt deleted file mode 100644 index 4e20334..0000000 --- a/ci/array-api-tests-rev.txt +++ /dev/null @@ -1 +0,0 @@ -c48410f96fc58e02eea844e6b7f6cc01680f77ce diff --git a/ci/clone_array_api_tests.sh b/ci/clone_array_api_tests.sh deleted file mode 100755 index c74e401..0000000 --- a/ci/clone_array_api_tests.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/usr/bin/env bash -set -euxo pipefail - -ARRAY_API_TESTS_DIR="${ARRAY_API_TESTS_DIR:-"../array-api-tests"}" -if [ ! -d "$ARRAY_API_TESTS_DIR" ]; then - git clone --recursive https://github.com/data-apis/array-api-tests.git "$ARRAY_API_TESTS_DIR" -fi - -git --git-dir="$ARRAY_API_TESTS_DIR/.git" --work-tree "$ARRAY_API_TESTS_DIR" clean -xddf -git --git-dir="$ARRAY_API_TESTS_DIR/.git" --work-tree "$ARRAY_API_TESTS_DIR" fetch -git --git-dir="$ARRAY_API_TESTS_DIR/.git" --work-tree "$ARRAY_API_TESTS_DIR" reset --hard $(cat "ci/array-api-tests-rev.txt") From ebca06345c382a4ba0ffa177da3af922f4b0305f Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 14:12:57 -0400 Subject: [PATCH 34/81] precommit --- tests/test_array_api.py | 78 +++++++++++++++++++++++++++++++---------- 1 file changed, 59 insertions(+), 19 deletions(-) diff --git a/tests/test_array_api.py b/tests/test_array_api.py index 77311f2..0a05bcf 100644 --- a/tests/test_array_api.py +++ b/tests/test_array_api.py @@ -4,9 +4,21 @@ def test_array_api(): - ARRAY_API_TESTS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), os.environ.get("ARRAY_API_TESTS_DIR", "../array-api-tests"))) - ARRAY_API_TESTS_REV = os.environ.get("ARRAY_API_TESTS_REV", "c48410f96fc58e02eea844e6b7f6cc01680f77ce") - ARRAY_API_TESTS_SKIPS = os.path.abspath(os.path.join(os.path.dirname(__file__), os.environ.get("ARRAY_API_TESTS_SKIPS", "array-api-skips.txt"))) + ARRAY_API_TESTS_DIR = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + os.environ.get("ARRAY_API_TESTS_DIR", "../array-api-tests"), + ) + ) + ARRAY_API_TESTS_REV = os.environ.get( + "ARRAY_API_TESTS_REV", "c48410f96fc58e02eea844e6b7f6cc01680f77ce" + ) + ARRAY_API_TESTS_SKIPS = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + os.environ.get("ARRAY_API_TESTS_SKIPS", "array-api-skips.txt"), + ) + ) ARRAY_API_TESTS_ARGS = os.environ.get("ARRAY_API_TESTS_ARGS", "-vv -s") print(f"[array-api] using dir: {ARRAY_API_TESTS_DIR}", flush=True) @@ -15,49 +27,77 @@ def test_array_api(): if not os.path.isdir(ARRAY_API_TESTS_DIR): print("[array-api] cloning test repo...", flush=True) subprocess.run( - ["git", "clone", "--recursive", - "https://github.com/data-apis/array-api-tests.git", - ARRAY_API_TESTS_DIR], - check=True + [ + "git", + "clone", + "--recursive", + "https://github.com/data-apis/array-api-tests.git", + ARRAY_API_TESTS_DIR, + ], + check=True, ) print("[array-api] cleaning repo...", flush=True) subprocess.run( - ["git", "--git-dir", f"{ARRAY_API_TESTS_DIR}/.git", - "--work-tree", ARRAY_API_TESTS_DIR, "clean", "-xddf"], - check=True + [ + "git", + "--git-dir", + f"{ARRAY_API_TESTS_DIR}/.git", + "--work-tree", + ARRAY_API_TESTS_DIR, + "clean", + "-xddf", + ], + check=True, ) print("[array-api] fetching latest refs...", flush=True) subprocess.run( - ["git", "--git-dir", f"{ARRAY_API_TESTS_DIR}/.git", - "--work-tree", ARRAY_API_TESTS_DIR, "fetch"], - check=True + [ + "git", + "--git-dir", + f"{ARRAY_API_TESTS_DIR}/.git", + "--work-tree", + ARRAY_API_TESTS_DIR, + "fetch", + ], + check=True, ) print("[array-api] checking out target revision...", flush=True) subprocess.run( - ["git", "--git-dir", f"{ARRAY_API_TESTS_DIR}/.git", - "--work-tree", ARRAY_API_TESTS_DIR, "reset", "--hard", ARRAY_API_TESTS_REV], - check=True + [ + "git", + "--git-dir", + f"{ARRAY_API_TESTS_DIR}/.git", + "--work-tree", + ARRAY_API_TESTS_DIR, + "reset", + "--hard", + ARRAY_API_TESTS_REV, + ], + check=True, ) # Run the tests using pytest print("[array-api] running external array-api-tests...", flush=True) result = subprocess.run( [ - sys.executable, "-m", "pytest", + sys.executable, + "-m", + "pytest", *ARRAY_API_TESTS_ARGS.split(), f"{ARRAY_API_TESTS_DIR}/array_api_tests/", "--max-examples=2", "--derandomize", "--disable-deadline", "--disable-warnings", - "--skips-file", ARRAY_API_TESTS_SKIPS + "--skips-file", + ARRAY_API_TESTS_SKIPS, ], env={**os.environ, "ARRAY_API_TESTS_MODULE": "finch", "PYTHONUNBUFFERED": "1"}, check=False, text=True, ) print("[array-api] array-api-tests completed!", flush=True) - assert result.returncode == 0, "Array API tests failed" \ No newline at end of file + assert result.returncode == 0, "Array API tests failed" From 1540f103fba9b94a8f536b7d55daf0f7c8084be4 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 14:33:43 -0400 Subject: [PATCH 35/81] learning --- pyproject.toml | 2 +- pytest.ini | 1 + src/finch/compiler.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d11a784..53cf4ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "numpy (>=1.19,<2.4)", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", - # "finch-tensor-lite (==0.3.0)", + "finch-tensor-lite (==0.3.0)", ] [tool.poetry] diff --git a/pytest.ini b/pytest.ini index a351985..89d89a7 100644 --- a/pytest.ini +++ b/pytest.ini @@ -4,5 +4,6 @@ filterwarnings = ignore::PendingDeprecationWarning testpaths = finch +norecursedirs = array-api-tests junit_family=xunit2 xfail_strict=true diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 433e7ee..8aa075f 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -3,7 +3,7 @@ import finchlite.finch_notation.nodes as ntn from finchlite.algebra import make_tuple -from finchlite.algebra.algebra import overwrite, promote_max, promote_min +from finchlite.algebra import overwrite, promote_max, promote_min from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary From d1697d9951b9d1971eae1ab93998f5f4f3d749e1 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 14:35:45 -0400 Subject: [PATCH 36/81] fix --- .gitignore | 2 +- pytest.ini | 1 + tests/test_array_api.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 2b2b89f..d9794b2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,4 @@ -array-api-tests +tests/array-api-tests # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/pytest.ini b/pytest.ini index a351985..981202e 100644 --- a/pytest.ini +++ b/pytest.ini @@ -4,5 +4,6 @@ filterwarnings = ignore::PendingDeprecationWarning testpaths = finch +norecursedirs = tests/array-api-tests junit_family=xunit2 xfail_strict=true diff --git a/tests/test_array_api.py b/tests/test_array_api.py index 0a05bcf..02c243a 100644 --- a/tests/test_array_api.py +++ b/tests/test_array_api.py @@ -7,7 +7,7 @@ def test_array_api(): ARRAY_API_TESTS_DIR = os.path.abspath( os.path.join( os.path.dirname(__file__), - os.environ.get("ARRAY_API_TESTS_DIR", "../array-api-tests"), + os.environ.get("ARRAY_API_TESTS_DIR", "array-api-tests"), ) ) ARRAY_API_TESTS_REV = os.environ.get( From 8af5440b7dfa11065601da8d32e9ec35767a6c4e Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 14:45:00 -0400 Subject: [PATCH 37/81] fix --- .gitignore | 2 +- ...array-api-skips.txt => array-api-skips.txt | 0 pytest.ini | 2 +- tests/test_array_api.py | 20 +++++++++---------- 4 files changed, 12 insertions(+), 12 deletions(-) rename tests/array-api-skips.txt => array-api-skips.txt (100%) diff --git a/.gitignore b/.gitignore index d9794b2..2b2b89f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,4 @@ -tests/array-api-tests +array-api-tests # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/tests/array-api-skips.txt b/array-api-skips.txt similarity index 100% rename from tests/array-api-skips.txt rename to array-api-skips.txt diff --git a/pytest.ini b/pytest.ini index 981202e..89d89a7 100644 --- a/pytest.ini +++ b/pytest.ini @@ -4,6 +4,6 @@ filterwarnings = ignore::PendingDeprecationWarning testpaths = finch -norecursedirs = tests/array-api-tests +norecursedirs = array-api-tests junit_family=xunit2 xfail_strict=true diff --git a/tests/test_array_api.py b/tests/test_array_api.py index 02c243a..b6440f1 100644 --- a/tests/test_array_api.py +++ b/tests/test_array_api.py @@ -4,20 +4,20 @@ def test_array_api(): - ARRAY_API_TESTS_DIR = os.path.abspath( - os.path.join( - os.path.dirname(__file__), - os.environ.get("ARRAY_API_TESTS_DIR", "array-api-tests"), - ) + ARRAY_API_TESTS_DIR = os.environ.get( + "ARRAY_API_TESTS_DIR", + os.path.abspath( + os.path.join(os.path.dirname(__file__), "../array-api-tests"), + ), ) ARRAY_API_TESTS_REV = os.environ.get( "ARRAY_API_TESTS_REV", "c48410f96fc58e02eea844e6b7f6cc01680f77ce" ) - ARRAY_API_TESTS_SKIPS = os.path.abspath( - os.path.join( - os.path.dirname(__file__), - os.environ.get("ARRAY_API_TESTS_SKIPS", "array-api-skips.txt"), - ) + ARRAY_API_TESTS_SKIPS = os.environ.get( + "ARRAY_API_TESTS_SKIPS", + os.path.abspath( + os.path.join(os.path.dirname(__file__), "../array-api-skips.txt"), + ), ) ARRAY_API_TESTS_ARGS = os.environ.get("ARRAY_API_TESTS_ARGS", "-vv -s") From 101e60b69395ca78e1d20512d3bbd9131b0d4a0b Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 14:46:12 -0400 Subject: [PATCH 38/81] cool --- src/finch/compiler.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 8aa075f..920e941 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -2,8 +2,7 @@ import operator import finchlite.finch_notation.nodes as ntn -from finchlite.algebra import make_tuple -from finchlite.algebra import overwrite, promote_max, promote_min +from finchlite.algebra import make_tuple, overwrite, promote_max, promote_min from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary From 836a016d93458ea4d42f1d9cec9f3fa60b5842cf Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 15:48:56 -0400 Subject: [PATCH 39/81] kinda messy but also correct --- src/finch/levels.py | 98 +++++++++++++++++++++++++++++++++++++++++---- src/finch/tensor.py | 97 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 7 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 5f85849..d78a127 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -75,11 +75,6 @@ def create_jl_obj(self) -> JuliaObj: class NestedLevelFType(LevelFType): - def __init__(self, lvl: LevelFType): - if not isinstance(lvl, NestedLevelFType | Element): - raise ValueError("lvl must be a NestedLevelFType or Element.") - self.lvl = lvl - @property def ndim(self) -> np.intp: return self.lvl.ndim + np.intp(1) @@ -102,21 +97,110 @@ def __hash__(self): def create_jl_obj(self) -> JuliaObj: ... +@dataclass(frozen=True) class Dense(NestedLevelFType): + lvl: NestedLevelFType + def create_jl_obj(self) -> JuliaObj: return jl.Dense(self.lvl.create_jl_obj()) - - + +@dataclass(frozen=True) class SparseList(NestedLevelFType): + lvl: NestedLevelFType + def create_jl_obj(self) -> JuliaObj: return jl.SparseList(self.lvl.create_jl_obj()) +@dataclass(frozen=True) +class SparseCOO(NestedLevelFType): + lvl: NestedLevelFType + N: int = 2 + def create_jl_obj(self) -> JuliaObj: + return jl.SparseCOO(self.lvl.create_jl_obj()) +@dataclass(frozen=True) class SparseByteMap(NestedLevelFType): + lvl: NestedLevelFType + def create_jl_obj(self) -> JuliaObj: return jl.SparseByteMap(self.lvl.create_jl_obj()) +class _Display: + _obj: JuliaObj + + def __repr__(self): + return jl.sprint(jl.show, self._obj) + + def __str__(self): + return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) + + +# LEVEL + + +class AbstractLevel(_Display): + pass + + +# core levels + + +class Dense(AbstractLevel): + def __init__(self, lvl, shape=None): + args = [lvl._obj] + if shape is not None: + args.append(shape) + self._obj = jl.Dense(*args) + + +class Element(AbstractLevel): + def __init__(self, fill_value, data=None): + args = [fill_value] + if data is not None: + args.append(data) + self._obj = jl.Element(*args) + + +class Pattern(AbstractLevel): + def __init__(self): + self._obj = jl.Pattern() + + +# advanced levels + + +class SparseList(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.SparseList(lvl._obj) + + +class SparseByteMap(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.SparseByteMap(lvl._obj) + +jl.PlusOneVector(arr) + +class RepeatRLE(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.RepeatRLE(lvl._obj) + + +class SparseVBL(AbstractLevel): + def __init__(self, lvl): + self._obj = jl.SparseVBL(lvl._obj) + + +class SparseCOO(AbstractLevel): + def __init__(self, ndim, lvl): + self._obj = jl.SparseCOO[ndim](lvl._obj) + + +class SparseHash(AbstractLevel): + def __init__(self, ndim, lvl): + self._obj = jl.SparseHash[ndim](lvl._obj) + + # Helper Methods def construct_levels(obj: JuliaObj, fill_value: number) -> LevelFType: if jl.isa(obj.lvl, jl.Finch.Element): diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 01311da..fa0ef22 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -132,3 +132,100 @@ def __repr__(self): def __str__(self): return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) + + def __array_namespace__(self, *, api_version: str | None = None) -> Any: + if api_version is None: + api_version = "2024.12" + + if api_version not in {"2021.12", "2022.12", "2023.12", "2024.12"}: + raise ValueError(f'"{api_version}" Array API version not supported.') + import finch + + return finch + +def asarray( + obj, + /, + *, + dtype: DType | None = None, + fill_value: np.number | None = None, + copy: bool | None = None, +) -> FinchJLTensor: + if fill_value is None: + fill_value = 0.0 + if isinstance(obj, FinchJLTensor): + if copy: + return obj.copy() + else: + return obj + elif isinstance(obj, np.ndarray): + if copy: + if np.isfortran(obj): + arr = arr.copy() + else: + obj = np.asfortranarray(obj) + dtype = arr.dtype.type + if ( + dtype == np.bool_ + ): # Fails with: Finch currently only supports isbits defaults + dtype = jl_dtypes.bool + lvl = ElementLevel(fill_value, arr.reshape(-1, order="F")) + for i in arr.shape: + lvl = DenseLevel(lvl, i) + return FinchJLTensor(lvl) + elif hasattr(x, "__module__") and x.__module__.startswith("scipy.sparse"): + if obj.format == "coo": + obj = obj.T + if copy: + if obj.format in ("coo", "csc"): + if not x.has_sorted_indices: + obj = obj.sorted_indices() + else: + obj = obj.copy() + if not x.has_canonical_format: + obj.sum_duplicates() + else: + obj = obj.asformat("csc") + if copy is False and not obj.format in ("coo", "csc") and not obj.has_canonical_format: + raise ValueError( + "Unable to avoid copy while creating an array as requested." + ) + m, n = obj.shape + if obj.format == "coo": + return Tensor( + SparseCOOLevel( + ElementLevel( + dtype, + fill_value, + obj.data + ), + 2, + idxs = ( + x.cols, + x.rows, + ), + ) + ) + elif x.format == "csc": + return Tensor( + DenseLevel( + SparseListLevel( + ElementLevel( + dtype, + fill_value, + obj.data + ), + n, + obj.indptr, + obj.indices + ), + (m, n) + ) + ) + else: + raise ValueError(f"Unsupported SciPy format: {type(x)}") + else: + raise ValueError( + "Either scalar, numpy, scipy.sparse or a raw julia object should " + f"be provided. Found: {type(obj)}" + ) From a3d58684c920066967177b9a8b35abfde3f8c1fc Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 16:12:20 -0400 Subject: [PATCH 40/81] move around --- src/finch/levels.py | 55 ++++++++++++++------------------------------- src/finch/scalar.py | 25 +++++++++++++++++++++ 2 files changed, 42 insertions(+), 38 deletions(-) create mode 100644 src/finch/scalar.py diff --git a/src/finch/levels.py b/src/finch/levels.py index d78a127..672f18e 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -7,6 +7,7 @@ from .julia import jl from .typing import JuliaObj, number +from dataclasses import dataclass # Abstract FTypes @@ -22,33 +23,33 @@ def __call__(self, _) -> Tensor: raise Exception("Cannot create an object of this type!") -class Scalar(LevelFType): - def __init__(self, val: number): - self._val = val +class NestedLevelFType(LevelFType): @property def ndim(self) -> np.intp: - return np.intp(0) + return self.lvl.ndim + np.intp(1) @property def fill_value(self) -> Any: - return self._val + return self.lvl.fill_value @property def element_type(self) -> Any: - return type(self._val) + return self.lvl.element_type def __eq__(self, other): - return isinstance(other, Scalar) and self._val == other._val + return type(other) is type(self) and self.lvl == other.lvl def __hash__(self): - return hash((self.__class__.__name__, self._val)) + return hash((self.__class__.__name__, self.lvl.__hash__)) + + @abstractmethod + def create_jl_obj(self) -> JuliaObj: ... - def create_jl_obj(self) -> JuliaObj: - return jl.Scalar(self._val) -class Element(LevelFType): + +class ElementFType(LevelFType): def __init__(self, fill_value: number): self._fill_value = fill_value @@ -65,7 +66,7 @@ def element_type(self) -> Any: return type(self._fill_value) def __eq__(self, other): - return isinstance(other, Element) and self._fill_value == other.fill_value + return isinstance(other, ElementFType) and self._fill_value == other.fill_value def __hash__(self): return hash((self.__class__.__name__, self._fill_value)) @@ -74,52 +75,30 @@ def create_jl_obj(self) -> JuliaObj: return jl.Element(self._fill_value) -class NestedLevelFType(LevelFType): - @property - def ndim(self) -> np.intp: - return self.lvl.ndim + np.intp(1) - - @property - def fill_value(self) -> Any: - return self.lvl.fill_value - - @property - def element_type(self) -> Any: - return self.lvl.element_type - - def __eq__(self, other): - return type(other) is type(self) and self.lvl == other.lvl - - def __hash__(self): - return hash((self.__class__.__name__, self.lvl.__hash__)) - - @abstractmethod - def create_jl_obj(self) -> JuliaObj: ... - @dataclass(frozen=True) -class Dense(NestedLevelFType): +class DenseFType(NestedLevelFType): lvl: NestedLevelFType def create_jl_obj(self) -> JuliaObj: return jl.Dense(self.lvl.create_jl_obj()) @dataclass(frozen=True) -class SparseList(NestedLevelFType): +class SparseListFType(NestedLevelFType): lvl: NestedLevelFType def create_jl_obj(self) -> JuliaObj: return jl.SparseList(self.lvl.create_jl_obj()) @dataclass(frozen=True) -class SparseCOO(NestedLevelFType): +class SparseCoOFType(NestedLevelFType): lvl: NestedLevelFType N: int = 2 def create_jl_obj(self) -> JuliaObj: return jl.SparseCOO(self.lvl.create_jl_obj()) @dataclass(frozen=True) -class SparseByteMap(NestedLevelFType): +class SparseByteMapFType(NestedLevelFType): lvl: NestedLevelFType def create_jl_obj(self) -> JuliaObj: diff --git a/src/finch/scalar.py b/src/finch/scalar.py new file mode 100644 index 0000000..47d4e8b --- /dev/null +++ b/src/finch/scalar.py @@ -0,0 +1,25 @@ + +class ScalarFType(LevelFType): + def __init__(self, val: number): + self._val = val + + @property + def ndim(self) -> np.intp: + return np.intp(0) + + @property + def fill_value(self) -> Any: + return self._val + + @property + def element_type(self) -> Any: + return type(self._val) + + def __eq__(self, other): + return isinstance(other, ScalarFType) and self._val == other._val + + def __hash__(self): + return hash((self.__class__.__name__, self._val)) + + def create_jl_obj(self) -> JuliaObj: + return jl.Scalar(self._val) \ No newline at end of file From 5f7e59bad5ca7d8c777e6f96966f18cd1f4c830f Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 17:14:39 -0400 Subject: [PATCH 41/81] good stuff --- src/finch/buffer.py | 71 +++++++++++ src/finch/levels.py | 291 +++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 346 insertions(+), 16 deletions(-) create mode 100644 src/finch/buffer.py diff --git a/src/finch/buffer.py b/src/finch/buffer.py new file mode 100644 index 0000000..429806f --- /dev/null +++ b/src/finch/buffer.py @@ -0,0 +1,71 @@ + +from abc import ABC +from finchlite import Buffer, NumpyBuffer, BufferFType +from .julia import jc, jl + +class PlusOneBufferFType(BufferFType): + def __init__(self, data_ftype): + self.data_ftype = data_ftype + + def __call__(self, *args, **kwargs): + return PlusOneBuffer(self.data_ftype(*args, **kwargs)) + + @property + def element_type(self): + return self.data_ftype.element_type + + @property + def length_type(self): + return self.data_ftype.length_type + + +class PlusOneBuffer(Buffer, ABC): + """ + Buffer that adds one to each element when loaded and subtracts one from each element when stored. + """ + + def __init__(self, data): + self.data : Buffer = data + + def ftype(self): + return PlusOneBufferFType(self.data.ftype()) + + def length(self): + return self.data.length() + + @property + def element_type(self): + """ + Return the type of elements stored in the buffer. + This is typically the same as the dtype used to create the buffer. + """ + return self.data.element_type + + @property + def length_type(self): + return self.data.length_type() + + def load(self, idx: int): + return self.data.load(idx) + 1 + + def store(self, idx: int, val): + self.data.store(idx, val - 1) + + def resize(self, len: int): + self.data.resize(len) + +def buffer_to_jlobj(buffer: Buffer): + if isinstance(buffer, PlusOneBuffer): + return jl.PlusOneVector(buffer_to_jlobj(buffer.data)) + elif isinstance(buffer, NumpyBuffer): + return buffer.arr + else: + raise ValueError(f"Unsupported buffer type: {type(buffer)}") + +def jlobj_to_buffer(jlobj): + if isinstance(jlobj, jl.PlusOneVector): + return PlusOneBuffer(jlobj_to_buffer(jlobj.data)) + elif isinstance(jlobj, np.ndarray): + return NumpyBuffer(jlobj) + else: + raise ValueError(f"Unsupported Julia object type: {type(jlobj)}") \ No newline at end of file diff --git a/src/finch/levels.py b/src/finch/levels.py index 672f18e..2c7fc26 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -47,9 +47,12 @@ def __hash__(self): def create_jl_obj(self) -> JuliaObj: ... - - class ElementFType(LevelFType): + """Element level storage format for scalar tensor leaves. + + A subfiber of an element level is a scalar, initialized to a fill value. + The element level is a leaf level used at the end of the tensor tree structure. + """ def __init__(self, fill_value: number): self._fill_value = fill_value @@ -78,6 +81,12 @@ def create_jl_obj(self) -> JuliaObj: @dataclass(frozen=True) class DenseFType(NestedLevelFType): + """Dense format wrapper type for Finch tensors. + + A dense format stores every slice of a tensor. A subfiber of a dense level + is an array which stores every slice A[:, ..., :, i] as a distinct subfiber. + Dense levels support both row-major and column-major access. + """ lvl: NestedLevelFType def create_jl_obj(self) -> JuliaObj: @@ -85,13 +94,26 @@ def create_jl_obj(self) -> JuliaObj: @dataclass(frozen=True) class SparseListFType(NestedLevelFType): + """Sparse list format wrapper type for Finch tensors. + + A sparse list format stores only potentially non-fill slices using a sorted list. + Slices that are entirely fill_value are omitted. This format is efficient for + tensors with sparse patterns and supports column-major reads and bulk updates. + """ lvl: NestedLevelFType def create_jl_obj(self) -> JuliaObj: return jl.SparseList(self.lvl.create_jl_obj()) @dataclass(frozen=True) -class SparseCoOFType(NestedLevelFType): +class SparseCOOFType(NestedLevelFType): + """Coordinate (COO) format wrapper type for Finch tensors. + + A coordinate format stores sparse tensors as lists of coordinates. + It uses N separate arrays to record which coordinates are stored, + with coordinates sorted in column-major order. This is a legacy format + maintained for backward compatibility. + """ lvl: NestedLevelFType N: int = 2 def create_jl_obj(self) -> JuliaObj: @@ -99,6 +121,12 @@ def create_jl_obj(self) -> JuliaObj: @dataclass(frozen=True) class SparseByteMapFType(NestedLevelFType): + """Sparse byte map format wrapper type for Finch tensors. + + A sparse byte map format uses a dense bitmap to encode which slices + are stored, similar to SparseList but supporting random access. + Only potentially non-fill slices are stored as subfibers. + """ lvl: NestedLevelFType def create_jl_obj(self) -> JuliaObj: @@ -126,6 +154,26 @@ class AbstractLevel(_Display): class Dense(AbstractLevel): + """Dense level storage format. + + A subfiber of a dense level is an array which stores every slice as a distinct + subfiber in the child level. Dense levels support efficient random access and + both column-major and out-of-order updates. + + Parameters + ---------- + lvl : AbstractLevel + The child level that will store each slice. + shape : int, optional + The size of the dimension at this level. If not provided, the shape + is inferred from the child level. + + Examples + -------- + Create a 2D dense tensor: + + >>> dense_2d = Dense(Dense(Element(0.0))) + """ def __init__(self, lvl, shape=None): args = [lvl._obj] if shape is not None: @@ -134,6 +182,26 @@ def __init__(self, lvl, shape=None): class Element(AbstractLevel): + """Element level storage format (leaf level). + + A subfiber of an element level is a scalar of a specified type, initialized + to a fill value. Element levels form the leaf nodes of the tensor tree and + store the actual data values. + + Parameters + ---------- + fill_value : float or int + The default fill value for elements that are not explicitly set. + data : array-like, optional + Optional vector to store the element data. If not provided, a new + vector is created. + + Examples + -------- + Create an element level with zero fill value: + + >>> elem = Element(0.0) + """ def __init__(self, fill_value, data=None): args = [fill_value] if data is not None: @@ -142,6 +210,19 @@ def __init__(self, fill_value, data=None): class Pattern(AbstractLevel): + """Pattern level storage format (leaf level). + + A subfiber of a pattern level is the boolean value true, but with a fill + value of false. Pattern levels are used to create tensors representing + which values are stored by other fibers, allowing for tracking sparsity + structure without storing values. + + Examples + -------- + Create a pattern level to track which elements are nonzero: + + >>> pattern = Pattern() + """ def __init__(self): self._obj = jl.Pattern() @@ -150,38 +231,216 @@ def __init__(self): class SparseList(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseList(lvl._obj) + """Sparse list level storage format. + + A subfiber of a sparse list level stores only potentially non-fill slices + using a sorted list to record which slices are stored. Slices that are + entirely fill_value are omitted. This format is efficient for sparse tensors + and supports column-major reads and bulk updates, but does not support + random access or random updates. + + Parameters + ---------- + lvl : AbstractLevel + The child level that will store each non-fill slice. + dim : int, optional + The size of the dimension at this level. If not provided, the dimension + is inferred from the child level. + ptr : array-like, optional + Array of positions/pointers for the sparse list. If not provided, + will be created internally. Converted to ndarray if provided. + idx : array-like, optional + Array of indices for the sparse list. If not provided, + will be created internally. Converted to ndarray if provided. + + Examples + -------- + Create a sparse matrix in CSC format: + + >>> sparse_matrix = SparseList(SparseList(Element(0.0))) + """ + def __init__(self, lvl, dim=None, ptr=None, idx=None): + args = [lvl._obj] + if dim is not None: + args.append(dim) + if ptr is not None: + args.append(np.asarray(ptr)) + if idx is not None: + args.append(np.asarray(idx)) + self._obj = jl.SparseList(*args) class SparseByteMap(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseByteMap(lvl._obj) + """Sparse byte map level storage format. + + Similar to SparseList, but uses a dense bitmap to encode which slices are + stored instead of a sparse index list. This allows for efficient random + access while still omitting fill_value slices. The byte map approach trades + memory for faster lookups. + + Parameters + ---------- + lvl : AbstractLevel + The child level that will store each non-fill slice. + dim : int, optional + The size of the dimension at this level. If not provided, the dimension + is inferred from the child level. + + Examples + -------- + Create a sparse matrix with byte map storage: + + >>> sparse_matrix = SparseByteMap(SparseByteMap(Element(0.0))) + """ + def __init__(self, lvl, dim=None): + args = [lvl._obj] + if dim is not None: + args.append(dim) + self._obj = jl.SparseByteMap(*args) -jl.PlusOneVector(arr) class RepeatRLE(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.RepeatRLE(lvl._obj) + """Run-length encoding level for repeated values. + + This level stores runs of repeated values, useful for data with + long sequences of identical values. + + Parameters + ---------- + lvl : AbstractLevel + The child level to store run information. + dim : int, optional + The size of the dimension at this level. If not provided, the dimension + is inferred from the child level. + """ + def __init__(self, lvl, dim=None): + args = [lvl._obj] + if dim is not None: + args.append(dim) + self._obj = jl.RepeatRLE(*args) class SparseVBL(AbstractLevel): - def __init__(self, lvl): - self._obj = jl.SparseVBL(lvl._obj) + """Sparse variable block list level storage format. + + Like SparseList, but stores contiguous slices together in blocks. + This can improve cache locality and performance for certain access patterns. + + Parameters + ---------- + lvl : AbstractLevel + The child level that will store each block of slices. + dim : int, optional + The size of the dimension at this level. If not provided, the dimension + is inferred from the child level. + """ + def __init__(self, lvl, dim=None): + args = [lvl._obj] + if dim is not None: + args.append(dim) + self._obj = jl.SparseVBL(*args) class SparseCOO(AbstractLevel): - def __init__(self, ndim, lvl): - self._obj = jl.SparseCOO[ndim](lvl._obj) + """Coordinate (COO) sparse format level storage. + + This level stores sparse data using N coordinate lists (one per dimension). + Coordinates are stored in column-major order. This is a legacy format + maintained for backward compatibility; consider using other sparse formats + for new code. + + Parameters + ---------- + ndim : int + Number of dimensions for the coordinate format. + lvl : AbstractLevel + The child level to store coordinate data. + dims : tuple of int, optional + Sizes of the last N dimensions. If not provided, dimensions are + inferred from the child level. + tbl : tuple of arrays, optional + A tuple of coordinate arrays (one per dimension). Each array stores + the coordinates for that dimension. Arrays are converted to ndarray + if provided. If not provided, will be created internally. + + Examples + -------- + Create a 2D sparse tensor in COO format: + + >>> coo_2d = SparseCOO(2, Element(0.0)) + """ + def __init__(self, ndim, lvl, dims=None, tbl=None): + args = [lvl._obj] + if dims is not None: + if isinstance(dims, (list, tuple)): + args.extend(dims) + else: + args.append(dims) + if tbl is not None: + if isinstance(tbl, (list, tuple)): + args.extend([np.asarray(t) for t in tbl]) + else: + args.append(np.asarray(tbl)) + self._obj = jl.SparseCOO[ndim](*args) class SparseHash(AbstractLevel): - def __init__(self, ndim, lvl): - self._obj = jl.SparseHash[ndim](lvl._obj) + """Hash table based sparse format level storage. + + Uses a hash table to store sparse data, supporting efficient random access + and random updates. This format is useful when you need flexible out-of-order + insertion of elements. + + Parameters + ---------- + ndim : int + Number of dimensions for the hash format. + lvl : AbstractLevel + The child level to store hash table data. + dims : tuple of int, optional + Sizes of the last N dimensions. If not provided, dimensions are + inferred from the child level. + + Examples + -------- + Create a 2D sparse tensor using hash storage: + + >>> hash_2d = SparseHash(2, Element(0.0)) + """ + def __init__(self, ndim, lvl, dims=None): + args = [lvl._obj] + if dims is not None: + if isinstance(dims, (list, tuple)): + args.extend(dims) + else: + args.append(dims) + self._obj = jl.SparseHash[ndim](*args) # Helper Methods def construct_levels(obj: JuliaObj, fill_value: number) -> LevelFType: + """Construct a level hierarchy from a Julia Finch tensor object. + + Recursively constructs a Python representation of the tensor's level structure + by inspecting the Julia object's levels. + + Parameters + ---------- + obj : JuliaObj + A Julia Finch tensor object whose levels will be inspected. + fill_value : float or int + The fill value used for the tensor's sparse representation. + + Returns + ------- + LevelFType + A Python representation of the level hierarchy. + + Raises + ------ + Exception + If an unsupported level type is encountered. + """ if jl.isa(obj.lvl, jl.Finch.Element): return Element(fill_value) if jl.isa(obj.lvl, jl.Finch.Dense): From ba41871cbea452d2fa853086557ae8498234dc2d Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Wed, 25 Mar 2026 17:19:39 -0400 Subject: [PATCH 42/81] fixed --- src/finch/levels.py | 40 ++++++++++++++++++++++++++++------------ 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 2c7fc26..301043c 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -3,7 +3,8 @@ import numpy as np -from finchlite import Tensor, TensorFType +from finchlite import Tensor, TensorFType, Buffer +from .buffer import PlusOneBuffer, NumpyBuffer, buffer_to_jlobj, jlobj_to_buffer from .julia import jl from .typing import JuliaObj, number @@ -202,12 +203,17 @@ class Element(AbstractLevel): >>> elem = Element(0.0) """ - def __init__(self, fill_value, data=None): + def __init__(self, fill_value, data: Buffer | None=None): args = [fill_value] if data is not None: - args.append(data) + args.append(buffer_to_jlobj(data)) self._obj = jl.Element(*args) + @property + def data(self) -> Buffer: + """Return the data buffer for this element level.""" + return jlobj_to_buffer(self._obj.data) + class Pattern(AbstractLevel): """Pattern level storage format (leaf level). @@ -263,12 +269,20 @@ def __init__(self, lvl, dim=None, ptr=None, idx=None): args = [lvl._obj] if dim is not None: args.append(dim) - if ptr is not None: - args.append(np.asarray(ptr)) - if idx is not None: - args.append(np.asarray(idx)) + if ptr is not None and idx is not None: + args.append(buffer_to_jlobj(ptr)) + args.append(buffer_to_jlobj(idx)) self._obj = jl.SparseList(*args) + @property + def ptr(self) -> Buffer: + """Return the coordinate buffers for this COO level.""" + return jlobj_to_buffer(self._obj.ptr) + + @property + def idx(self) -> Buffer: + """Return the coordinate buffers for this COO level.""" + return jlobj_to_buffer(self._obj.idx) class SparseByteMap(AbstractLevel): """Sparse byte map level storage format. @@ -369,7 +383,7 @@ class SparseCOO(AbstractLevel): >>> coo_2d = SparseCOO(2, Element(0.0)) """ - def __init__(self, ndim, lvl, dims=None, tbl=None): + def __init__(self, ndim, lvl, dims=None, tbl : tuple[Buffer] | None =None): args = [lvl._obj] if dims is not None: if isinstance(dims, (list, tuple)): @@ -377,11 +391,13 @@ def __init__(self, ndim, lvl, dims=None, tbl=None): else: args.append(dims) if tbl is not None: - if isinstance(tbl, (list, tuple)): - args.extend([np.asarray(t) for t in tbl]) - else: - args.append(np.asarray(tbl)) + args.extend([buffer_to_jlobj(t) for t in tbl]) self._obj = jl.SparseCOO[ndim](*args) + + @property + def tbl(self) -> tuple[Buffer, ...]: + """Return the coordinate buffers for this COO level.""" + return tuple(jlobj_to_buffer(coord) for coord in self._obj.tbl) class SparseHash(AbstractLevel): From d678ad324ba21adaa4ca1d8dc8a46b5b3c986fe6 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 15:06:36 -0400 Subject: [PATCH 43/81] solid --- src/finch/levels.py | 6 +- tests/test_asarray.py | 347 ++++++++++++++++++++++++++++++++++++++++++ tests/test_levels.py | 346 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 696 insertions(+), 3 deletions(-) create mode 100644 tests/test_asarray.py create mode 100644 tests/test_levels.py diff --git a/src/finch/levels.py b/src/finch/levels.py index 301043c..9098e82 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -276,12 +276,12 @@ def __init__(self, lvl, dim=None, ptr=None, idx=None): @property def ptr(self) -> Buffer: - """Return the coordinate buffers for this COO level.""" + """Return the pointer array for this sparse list level.""" return jlobj_to_buffer(self._obj.ptr) @property def idx(self) -> Buffer: - """Return the coordinate buffers for this COO level.""" + """Return the index array for this sparse list level.""" return jlobj_to_buffer(self._obj.idx) class SparseByteMap(AbstractLevel): @@ -396,7 +396,7 @@ def __init__(self, ndim, lvl, dims=None, tbl : tuple[Buffer] | None =None): @property def tbl(self) -> tuple[Buffer, ...]: - """Return the coordinate buffers for this COO level.""" + """Return the coordinate array tuple for this COO level.""" return tuple(jlobj_to_buffer(coord) for coord in self._obj.tbl) diff --git a/tests/test_asarray.py b/tests/test_asarray.py new file mode 100644 index 0000000..7b1f6f3 --- /dev/null +++ b/tests/test_asarray.py @@ -0,0 +1,347 @@ +"""Tests for the asarray function.""" + +import pytest +import numpy as np +from finch import asarray +from finch.tensor import FinchJLTensor + +try: + import scipy.sparse as sp + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + + +class TestAsarrayNumpy: + """Test asarray with numpy arrays.""" + + def test_asarray_1d_array(self): + """Test converting a 1D numpy array.""" + arr = np.array([1.0, 2.0, 3.0]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_2d_array(self): + """Test converting a 2D numpy array.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_3d_array(self): + """Test converting a 3D numpy array.""" + arr = np.arange(24).reshape(2, 3, 4).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_fortran_order(self): + """Test converting a Fortran-order array.""" + arr = np.asfortranarray([[1.0, 2.0], [3.0, 4.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_with_fill_value(self): + """Test asarray with explicit fill_value.""" + arr = np.array([[1.0, 0.0], [0.0, 2.0]]) + result = asarray(arr, fill_value=0.0) + assert isinstance(result, FinchJLTensor) + + def test_asarray_with_dtype(self): + """Test asarray with explicit dtype.""" + arr = np.array([[1, 2], [3, 4]], dtype=np.int32) + result = asarray(arr, dtype=np.int32) + assert isinstance(result, FinchJLTensor) + + def test_asarray_default_fill_value(self): + """Test that default fill_value is 0.0.""" + arr = np.array([1.0, 2.0, 3.0]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_large_array(self): + """Test converting a large numpy array.""" + arr = np.arange(1000).reshape(10, 100).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_zero_array(self): + """Test converting an all-zero array.""" + arr = np.zeros((5, 5)) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_ones_array(self): + """Test converting an all-ones array.""" + arr = np.ones((5, 5)) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + +class TestAsarrayFinchTensor: + """Test asarray with FinchJLTensor inputs.""" + + def test_asarray_finch_tensor_no_copy(self): + """Test asarray on FinchJLTensor returns same object when copy=False.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + tensor1 = asarray(arr) + tensor2 = asarray(tensor1, copy=False) + assert tensor1 is tensor2 + + def test_asarray_finch_tensor_with_copy(self): + """Test asarray on FinchJLTensor with copy=True.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + tensor1 = asarray(arr) + tensor2 = asarray(tensor1, copy=True) + # Should be different objects + assert tensor1 is not tensor2 + assert isinstance(tensor2, FinchJLTensor) + + def test_asarray_finch_tensor_default_copy(self): + """Test asarray on FinchJLTensor with default copy behavior.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + tensor1 = asarray(arr) + tensor2 = asarray(tensor1) + # Default should return same object (copy=None defaults to no copy) + assert tensor1 is tensor2 + + +class TestAsarrayEdgeCases: + """Test asarray edge cases and special values.""" + + def test_asarray_single_element(self): + """Test converting single element array.""" + arr = np.array([5.0]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_negative_values(self): + """Test converting array with negative values.""" + arr = np.array([[-1.0, -2.0], [3.0, 4.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_mixed_values(self): + """Test converting array with mixed positive and negative values.""" + arr = np.array([[-5.0, 0.0, 5.0], [1.0, -1.0, 2.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_very_small_values(self): + """Test converting array with very small values.""" + arr = np.array([[1e-10, 1e-15], [1e-20, 1e-25]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_very_large_values(self): + """Test converting array with very large values.""" + arr = np.array([[1e10, 1e15], [1e20, 1e25]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_nan_values(self): + """Test converting array with NaN values.""" + arr = np.array([[1.0, np.nan], [np.nan, 4.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_inf_values(self): + """Test converting array with infinity values.""" + arr = np.array([[1.0, np.inf], [-np.inf, 4.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + +class TestAsarrayDataTypes: + """Test asarray with different data types.""" + + def test_asarray_float32(self): + """Test converting float32 array.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_float64(self): + """Test converting float64 array.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_int32(self): + """Test converting int32 array.""" + arr = np.array([[1, 2], [3, 4]], dtype=np.int32) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_int64(self): + """Test converting int64 array.""" + arr = np.array([[1, 2], [3, 4]], dtype=np.int64) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_complex64(self): + """Test converting complex64 array.""" + arr = np.array([[1+2j, 3+4j], [5+6j, 7+8j]], dtype=np.complex64) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_complex128(self): + """Test converting complex128 array.""" + arr = np.array([[1+2j, 3+4j], [5+6j, 7+8j]], dtype=np.complex128) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + +class TestAsarrayShapes: + """Test asarray with various array shapes.""" + + def test_asarray_row_vector(self): + """Test converting row vector.""" + arr = np.array([[1.0, 2.0, 3.0, 4.0, 5.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_column_vector(self): + """Test converting column vector.""" + arr = np.array([[1.0], [2.0], [3.0], [4.0], [5.0]]) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_square_matrix(self): + """Test converting square matrix.""" + arr = np.arange(25).reshape(5, 5).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_rectangular_matrix(self): + """Test converting rectangular matrix.""" + arr = np.arange(20).reshape(4, 5).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_tall_matrix(self): + """Test converting tall matrix.""" + arr = np.arange(20).reshape(10, 2).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_wide_matrix(self): + """Test converting wide matrix.""" + arr = np.arange(20).reshape(2, 10).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_4d_array(self): + """Test converting 4D array.""" + arr = np.arange(120).reshape(2, 3, 4, 5).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + def test_asarray_5d_array(self): + """Test converting 5D array.""" + arr = np.arange(120).reshape(2, 3, 4, 5, 1).astype(float) + result = asarray(arr) + assert isinstance(result, FinchJLTensor) + + +@pytest.mark.skipif(not HAS_SCIPY, reason="scipy not installed") +class TestAsarrayScipy: + """Test asarray with scipy.sparse matrices.""" + + def test_asarray_scipy_csc(self): + """Test converting scipy CSC sparse matrix.""" + data = np.array([1.0, 2.0, 3.0, 4.0]) + row = np.array([0, 1, 0, 2]) + col = np.array([0, 1, 1, 2]) + csc_matrix = sp.csr_matrix((data, (row, col)), shape=(3, 3)).tocsc() + result = asarray(csc_matrix) + assert isinstance(result, FinchJLTensor) + + def test_asarray_scipy_coo(self): + """Test converting scipy COO sparse matrix.""" + data = np.array([1.0, 2.0, 3.0, 4.0]) + row = np.array([0, 1, 0, 2]) + col = np.array([0, 1, 1, 2]) + coo_matrix = sp.coo_matrix((data, (row, col)), shape=(3, 3)) + result = asarray(coo_matrix) + assert isinstance(result, FinchJLTensor) + + def test_asarray_scipy_dense_csc(self): + """Test converting dense scipy matrix in CSC format.""" + arr = np.array([[1.0, 0.0, 2.0], [0.0, 3.0, 0.0], [4.0, 0.0, 5.0]]) + csc_matrix = sp.csc_matrix(arr) + result = asarray(csc_matrix) + assert isinstance(result, FinchJLTensor) + + def test_asarray_scipy_with_fill_value(self): + """Test asarray on scipy matrix with fill_value.""" + data = np.array([1.0, 2.0, 3.0]) + row = np.array([0, 1, 2]) + col = np.array([0, 1, 2]) + csc_matrix = sp.csr_matrix((data, (row, col)), shape=(3, 3)).tocsc() + result = asarray(csc_matrix, fill_value=0.0) + assert isinstance(result, FinchJLTensor) + + def test_asarray_scipy_with_copy_true(self): + """Test asarray on scipy matrix with copy=True.""" + data = np.array([1.0, 2.0, 3.0]) + row = np.array([0, 1, 2]) + col = np.array([0, 1, 2]) + csc_matrix = sp.csr_matrix((data, (row, col)), shape=(3, 3)).tocsc() + result = asarray(csc_matrix, copy=True) + assert isinstance(result, FinchJLTensor) + + def test_asarray_scipy_sorted_indices(self): + """Test asarray on scipy matrix with sorted indices.""" + data = np.array([1.0, 2.0, 3.0, 4.0]) + row = np.array([0, 1, 0, 2]) + col = np.array([0, 1, 1, 2]) + coo_matrix = sp.coo_matrix((data, (row, col)), shape=(3, 3)) + csc_matrix = coo_matrix.tocsc() + csc_matrix.sort_indices() + result = asarray(csc_matrix) + assert isinstance(result, FinchJLTensor) + + +class TestAsarrayErrors: + """Test asarray error handling.""" + + def test_asarray_invalid_type(self): + """Test asarray with unsupported type.""" + with pytest.raises((ValueError, TypeError, AttributeError)): + asarray("invalid string input") + + def test_asarray_invalid_list(self): + """Test asarray with plain Python list (should fail).""" + with pytest.raises((ValueError, TypeError, AttributeError)): + asarray([1, 2, 3]) + + def test_asarray_dict_input(self): + """Test asarray with dict input.""" + with pytest.raises((ValueError, TypeError, AttributeError)): + asarray({"a": 1}) + + def test_asarray_none_input(self): + """Test asarray with None input.""" + with pytest.raises((ValueError, TypeError, AttributeError)): + asarray(None) + + +class TestAsarrayOptions: + """Test asarray option combinations.""" + + def test_asarray_copy_none(self): + """Test asarray with copy=None.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + result = asarray(arr, copy=None) + assert isinstance(result, FinchJLTensor) + + def test_asarray_all_options(self): + """Test asarray with all options specified.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + result = asarray(arr, dtype=np.float64, fill_value=0.0, copy=True) + assert isinstance(result, FinchJLTensor) + + def test_asarray_numpy_no_copy(self): + """Test asarray on numpy with copy=False.""" + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + result = asarray(arr, copy=False) + assert isinstance(result, FinchJLTensor) diff --git a/tests/test_levels.py b/tests/test_levels.py new file mode 100644 index 0000000..80e8f04 --- /dev/null +++ b/tests/test_levels.py @@ -0,0 +1,346 @@ +"""Tests for the Finch levels module.""" + +import pytest +import numpy as np +from finch.levels import ( + Element, + Dense, + Pattern, + SparseList, + SparseByteMap, + RepeatRLE, + SparseVBL, + SparseCOO, + SparseHash, +) +from finch.buffer import NumpyBuffer + + +class TestElement: + """Test Element level construction.""" + + def test_element_creation_basic(self): + """Test creating an Element level with a fill value.""" + elem = Element(0.0) + assert elem._obj is not None + + def test_element_creation_with_int_fill(self): + """Test creating an Element level with integer fill value.""" + elem = Element(0) + assert elem._obj is not None + + def test_element_creation_with_float_fill(self): + """Test creating an Element level with float fill value.""" + elem = Element(3.14) + assert elem._obj is not None + + +class TestDense: + """Test Dense level construction.""" + + def test_dense_creation_basic(self): + """Test creating a Dense level.""" + elem = Element(0.0) + dense = Dense(elem) + assert dense._obj is not None + + def test_dense_creation_with_shape(self): + """Test creating a Dense level with explicit shape.""" + elem = Element(0.0) + dense = Dense(elem, shape=10) + assert dense._obj is not None + + def test_dense_nesting(self): + """Test creating nested Dense levels.""" + elem = Element(0.0) + dense1 = Dense(elem) + dense2 = Dense(dense1) + assert dense2._obj is not None + + +class TestPattern: + """Test Pattern level construction.""" + + def test_pattern_creation(self): + """Test creating a Pattern level.""" + pattern = Pattern() + assert pattern._obj is not None + + +class TestSparseList: + """Test SparseList level construction.""" + + def test_sparselist_creation_basic(self): + """Test creating a SparseList level.""" + elem = Element(0.0) + sparse = SparseList(elem) + assert sparse._obj is not None + + def test_sparselist_creation_with_dim(self): + """Test creating a SparseList level with explicit dimension.""" + elem = Element(0.0) + sparse = SparseList(elem, dim=10) + assert sparse._obj is not None + + def test_sparselist_creation_with_data_arrays(self): + """Test creating a SparseList level with pointer and index arrays.""" + elem = Element(0.0) + ptr = NumpyBuffer(np.array([0, 2, 2, 4], dtype=np.int32)) + idx = NumpyBuffer(np.array([1, 2, 1, 3], dtype=np.int32)) + sparse = SparseList(elem, ptr=ptr, idx=idx) + assert sparse._obj is not None + + def test_sparselist_creation_with_data_lists(self): + """Test creating a SparseList level with pointer and index as lists.""" + elem = Element(0.0) + ptr = [0, 2, 2, 4] + idx = [1, 2, 1, 3] + sparse = SparseList(elem, ptr=ptr, idx=idx) + assert sparse._obj is not None + + def test_sparselist_ptr_property(self): + """Test accessing ptr property of SparseList.""" + elem = Element(0.0) + sparse = SparseList(elem) + # Property should be accessible + ptr_buffer = sparse.ptr + assert ptr_buffer is not None + + def test_sparselist_idx_property(self): + """Test accessing idx property of SparseList.""" + elem = Element(0.0) + sparse = SparseList(elem) + # Property should be accessible + idx_buffer = sparse.idx + assert idx_buffer is not None + + +class TestSparseByteMap: + """Test SparseByteMap level construction.""" + + def test_sparsebytemap_creation_basic(self): + """Test creating a SparseByteMap level.""" + elem = Element(0.0) + sparse = SparseByteMap(elem) + assert sparse._obj is not None + + def test_sparsebytemap_creation_with_dim(self): + """Test creating a SparseByteMap level with explicit dimension.""" + elem = Element(0.0) + sparse = SparseByteMap(elem, dim=10) + assert sparse._obj is not None + + +class TestRepeatRLE: + """Test RepeatRLE level construction.""" + + def test_repeatrle_creation_basic(self): + """Test creating a RepeatRLE level.""" + elem = Element(0.0) + rle = RepeatRLE(elem) + assert rle._obj is not None + + def test_repeatrle_creation_with_dim(self): + """Test creating a RepeatRLE level with explicit dimension.""" + elem = Element(0.0) + rle = RepeatRLE(elem, dim=10) + assert rle._obj is not None + + +class TestSparseVBL: + """Test SparseVBL level construction.""" + + def test_sparsevbl_creation_basic(self): + """Test creating a SparseVBL level.""" + elem = Element(0.0) + vbl = SparseVBL(elem) + assert vbl._obj is not None + + def test_sparsevbl_creation_with_dim(self): + """Test creating a SparseVBL level with explicit dimension.""" + elem = Element(0.0) + vbl = SparseVBL(elem, dim=10) + assert vbl._obj is not None + + +class TestSparseCOO: + """Test SparseCOO level construction.""" + + def test_sparsecoo_creation_basic(self): + """Test creating a SparseCOO level.""" + elem = Element(0.0) + coo = SparseCOO(2, elem) + assert coo._obj is not None + + def test_sparsecoo_creation_with_dims(self): + """Test creating a SparseCOO level with explicit dimensions.""" + elem = Element(0.0) + coo = SparseCOO(2, elem, dims=(4, 3)) + assert coo._obj is not None + + def test_sparsecoo_creation_with_dims_list(self): + """Test creating a SparseCOO level with dimensions as list.""" + elem = Element(0.0) + coo = SparseCOO(2, elem, dims=[4, 3]) + assert coo._obj is not None + + def test_sparsecoo_creation_with_coordinate_arrays(self): + """Test creating a SparseCOO level with coordinate arrays.""" + elem = Element(0.0) + i_coords = NumpyBuffer(np.array([0, 1, 2, 3], dtype=np.int32)) + j_coords = NumpyBuffer(np.array([0, 0, 2, 2], dtype=np.int32)) + coo = SparseCOO(2, elem, tbl=(i_coords, j_coords)) + assert coo._obj is not None + + def test_sparsecoo_creation_with_coordinate_lists(self): + """Test creating a SparseCOO level with coordinate arrays as lists.""" + elem = Element(0.0) + i_coords = [0, 1, 2, 3] + j_coords = [0, 0, 2, 2] + coo = SparseCOO(2, elem, tbl=(i_coords, j_coords)) + assert coo._obj is not None + + def test_sparsecoo_3d(self): + """Test creating a 3D SparseCOO level.""" + elem = Element(0.0) + coo = SparseCOO(3, elem, dims=(5, 4, 3)) + assert coo._obj is not None + + def test_sparsecoo_tbl_property(self): + """Test accessing tbl property of SparseCOO.""" + elem = Element(0.0) + coo = SparseCOO(2, elem) + # Property should be accessible + tbl = coo.tbl + assert tbl is not None + assert isinstance(tbl, tuple) + + +class TestSparseHash: + """Test SparseHash level construction.""" + + def test_sparsehash_creation_basic(self): + """Test creating a SparseHash level.""" + elem = Element(0.0) + hash_level = SparseHash(2, elem) + assert hash_level._obj is not None + + def test_sparsehash_creation_with_dims(self): + """Test creating a SparseHash level with explicit dimensions.""" + elem = Element(0.0) + hash_level = SparseHash(2, elem, dims=(4, 3)) + assert hash_level._obj is not None + + def test_sparsehash_creation_with_dims_list(self): + """Test creating a SparseHash level with dimensions as list.""" + elem = Element(0.0) + hash_level = SparseHash(2, elem, dims=[4, 3]) + assert hash_level._obj is not None + + def test_sparsehash_3d(self): + """Test creating a 3D SparseHash level.""" + elem = Element(0.0) + hash_level = SparseHash(3, elem, dims=(5, 4, 3)) + assert hash_level._obj is not None + + +class TestComposedLevels: + """Test composed level hierarchies.""" + + def test_csc_matrix_format(self): + """Test creating CSC matrix format (Dense(SparseList(Element))).""" + elem = Element(0.0) + sparse = SparseList(elem) + dense = Dense(sparse) + assert dense._obj is not None + + def test_csr_like_format(self): + """Test creating CSR-like format (SparseList(Dense(Element))).""" + elem = Element(0.0) + dense = Dense(elem) + sparse = SparseList(dense) + assert sparse._obj is not None + + def test_dcsc_format(self): + """Test creating DCSC format (SparseList(SparseList(Element))).""" + elem = Element(0.0) + sparse1 = SparseList(elem) + sparse2 = SparseList(sparse1) + assert sparse2._obj is not None + + def test_deep_nesting(self): + """Test deeply nested levels.""" + elem = Element(0.0) + sparse = SparseList(elem) + dense = Dense(sparse) + sparse2 = SparseList(dense) + dense2 = Dense(sparse2) + assert dense2._obj is not None + + +class TestEdgeCases: + """Test edge cases and error conditions.""" + + def test_element_with_negative_fill(self): + """Test Element with negative fill value.""" + elem = Element(-1.0) + assert elem._obj is not None + + def test_sparselist_only_ptr_no_idx(self): + """Test SparseList with ptr but no idx (should not add both).""" + elem = Element(0.0) + ptr = NumpyBuffer(np.array([0, 2, 2], dtype=np.int32)) + sparse = SparseList(elem, ptr=ptr) + assert sparse._obj is not None + + def test_sparselist_only_idx_no_ptr(self): + """Test SparseList with idx but no ptr (should not add both).""" + elem = Element(0.0) + idx = NumpyBuffer(np.array([1, 2], dtype=np.int32)) + sparse = SparseList(elem, idx=idx) + assert sparse._obj is not None + + def test_sparsecoo_single_coordinate(self): + """Test SparseCOO with single coordinate.""" + elem = Element(0.0) + coords = NumpyBuffer(np.array([0], dtype=np.int32)) + coo = SparseCOO(1, elem, tbl=(coords,)) + assert coo._obj is not None + + def test_large_dimension(self): + """Test levels with large dimensions.""" + elem = Element(0.0) + sparse = SparseList(elem, dim=1000000) + assert sparse._obj is not None + + +class TestArrayConversion: + """Test that array arguments are properly converted.""" + + def test_sparselist_converts_lists_to_arrays(self): + """Test that SparseList converts list arguments to arrays.""" + elem = Element(0.0) + ptr = NumpyBuffer(np.array([0, 2, 4], dtype=np.int32)) + idx = NumpyBuffer(np.array([1, 2, 3], dtype=np.int32)) + sparse = SparseList(elem, ptr=ptr, idx=idx) + # Should not raise an error during creation + assert sparse._obj is not None + + def test_sparsecoo_converts_lists_to_arrays(self): + """Test that SparseCOO converts list arguments to arrays.""" + elem = Element(0.0) + coords_list = ( + NumpyBuffer(np.array([0, 1, 2], dtype=np.int32)), + NumpyBuffer(np.array([0, 1, 2], dtype=np.int32)) + ) + coo = SparseCOO(2, elem, tbl=coords_list) + # Should not raise an error during creation + assert coo._obj is not None + + def test_different_dtype_arrays(self): + """Test that different dtype arrays are handled.""" + elem = Element(0.0) + ptr = NumpyBuffer(np.array([0, 2, 4], dtype=np.int64)) + idx = NumpyBuffer(np.array([1, 2, 3], dtype=np.int32)) + sparse = SparseList(elem, ptr=ptr, idx=idx) + assert sparse._obj is not None From e99f363bcf4f744a65212ce8911d5a288afc8ad4 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 18:28:54 -0400 Subject: [PATCH 44/81] fix imports --- src/finch/buffer.py | 3 ++- src/finch/levels.py | 3 ++- src/finch/tensor.py | 3 +-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/finch/buffer.py b/src/finch/buffer.py index 429806f..664cffe 100644 --- a/src/finch/buffer.py +++ b/src/finch/buffer.py @@ -1,6 +1,7 @@ from abc import ABC -from finchlite import Buffer, NumpyBuffer, BufferFType +from finchlite.finch_assembly import Buffer, BufferFType +from finchlite.codegen import NumpyBuffer from .julia import jc, jl class PlusOneBufferFType(BufferFType): diff --git a/src/finch/levels.py b/src/finch/levels.py index 9098e82..a1703fa 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -3,7 +3,8 @@ import numpy as np -from finchlite import Tensor, TensorFType, Buffer +from finchlite import Tensor, TensorFType +from finchlite.finch_assembly import Buffer, BufferFType from .buffer import PlusOneBuffer, NumpyBuffer, buffer_to_jlobj, jlobj_to_buffer from .julia import jl diff --git a/src/finch/tensor.py b/src/finch/tensor.py index fa0ef22..ef2da6c 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -5,11 +5,10 @@ from finchlite import EagerTensor, Tensor, TensorFType from .julia import jc, jl -from .levels import LevelFType, Scalar, construct_levels +from .levels import LevelFType, construct_levels from .typing import JuliaObj from .utils import add_missing_dims, add_plus_one, expand_ellipsis - # Tensor Class and associated ftype class FinchJLTensorFType(TensorFType): def __init__(self, lvl): From d0b43043c7d996cd5fe313a3e3e53e8c26ca4fe2 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 19:42:58 -0400 Subject: [PATCH 45/81] fix --- pytest.ini | 2 +- src/finch/__init__.py | 1 + src/finch/julia.py | 2 + src/finch/levels.py | 40 ++++----- src/finch/scheduler.py | 6 +- src/finch/tensor.py | 31 +++---- tests/conftest.py | 1 - tests/test_compiler.py | 4 +- tests/test_levels.py | 190 ++++++++++++++++++++--------------------- 9 files changed, 140 insertions(+), 137 deletions(-) diff --git a/pytest.ini b/pytest.ini index 89d89a7..e00b2ce 100644 --- a/pytest.ini +++ b/pytest.ini @@ -3,7 +3,7 @@ addopts = --cov-report term-missing --cov-report html --cov-report=xml --cov-rep filterwarnings = ignore::PendingDeprecationWarning testpaths = - finch + tests norecursedirs = array-api-tests junit_family=xunit2 xfail_strict=true diff --git a/src/finch/__init__.py b/src/finch/__init__.py index 40b13e9..643a54d 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -122,6 +122,7 @@ ) from .scheduler import COMPILE_JULIA from .tensor import ( + asarray, FinchJLTensor, FinchJLTensorFType, ) diff --git a/src/finch/julia.py b/src/finch/julia.py index c8f2191..f5f087a 100644 --- a/src/finch/julia.py +++ b/src/finch/julia.py @@ -1,3 +1,5 @@ +import os # noqa: I001, F401 +os.environ["PYTHON_JULIACALL_HANDLE_SIGNALS"] = "yes" import juliapkg # noqa: I001, F401 # To change the version of Finch used, see the documentation for pyjuliapkg here: https://github.com/JuliaPy/pyjuliapkg diff --git a/src/finch/levels.py b/src/finch/levels.py index a1703fa..705526d 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -155,7 +155,7 @@ class AbstractLevel(_Display): # core levels -class Dense(AbstractLevel): +class DenseLevel(AbstractLevel): """Dense level storage format. A subfiber of a dense level is an array which stores every slice as a distinct @@ -174,7 +174,7 @@ class Dense(AbstractLevel): -------- Create a 2D dense tensor: - >>> dense_2d = Dense(Dense(Element(0.0))) + >>> dense_2d = DenseLevel(DenseLevel(ElementLevel(0.0))) """ def __init__(self, lvl, shape=None): args = [lvl._obj] @@ -183,7 +183,7 @@ def __init__(self, lvl, shape=None): self._obj = jl.Dense(*args) -class Element(AbstractLevel): +class ElementLevel(AbstractLevel): """Element level storage format (leaf level). A subfiber of an element level is a scalar of a specified type, initialized @@ -202,7 +202,7 @@ class Element(AbstractLevel): -------- Create an element level with zero fill value: - >>> elem = Element(0.0) + >>> elem = ElementLevel(0.0) """ def __init__(self, fill_value, data: Buffer | None=None): args = [fill_value] @@ -216,7 +216,7 @@ def data(self) -> Buffer: return jlobj_to_buffer(self._obj.data) -class Pattern(AbstractLevel): +class PatternLevel(AbstractLevel): """Pattern level storage format (leaf level). A subfiber of a pattern level is the boolean value true, but with a fill @@ -228,7 +228,7 @@ class Pattern(AbstractLevel): -------- Create a pattern level to track which elements are nonzero: - >>> pattern = Pattern() + >>> pattern = PatternLevel() """ def __init__(self): self._obj = jl.Pattern() @@ -237,7 +237,7 @@ def __init__(self): # advanced levels -class SparseList(AbstractLevel): +class SparseListLevel(AbstractLevel): """Sparse list level storage format. A subfiber of a sparse list level stores only potentially non-fill slices @@ -264,7 +264,7 @@ class SparseList(AbstractLevel): -------- Create a sparse matrix in CSC format: - >>> sparse_matrix = SparseList(SparseList(Element(0.0))) + >>> sparse_matrix = SparseListLevel(SparseListLevel(ElementLevel(0.0))) """ def __init__(self, lvl, dim=None, ptr=None, idx=None): args = [lvl._obj] @@ -285,7 +285,7 @@ def idx(self) -> Buffer: """Return the index array for this sparse list level.""" return jlobj_to_buffer(self._obj.idx) -class SparseByteMap(AbstractLevel): +class SparseByteMapLevel(AbstractLevel): """Sparse byte map level storage format. Similar to SparseList, but uses a dense bitmap to encode which slices are @@ -305,7 +305,7 @@ class SparseByteMap(AbstractLevel): -------- Create a sparse matrix with byte map storage: - >>> sparse_matrix = SparseByteMap(SparseByteMap(Element(0.0))) + >>> sparse_matrix = SparseByteMapLevel(SparseByteMapLevel(ElementLevel(0.0))) """ def __init__(self, lvl, dim=None): args = [lvl._obj] @@ -314,7 +314,7 @@ def __init__(self, lvl, dim=None): self._obj = jl.SparseByteMap(*args) -class RepeatRLE(AbstractLevel): +class RepeatRLELevel(AbstractLevel): """Run-length encoding level for repeated values. This level stores runs of repeated values, useful for data with @@ -335,7 +335,7 @@ def __init__(self, lvl, dim=None): self._obj = jl.RepeatRLE(*args) -class SparseVBL(AbstractLevel): +class SparseVBLLevel(AbstractLevel): """Sparse variable block list level storage format. Like SparseList, but stores contiguous slices together in blocks. @@ -356,7 +356,7 @@ def __init__(self, lvl, dim=None): self._obj = jl.SparseVBL(*args) -class SparseCOO(AbstractLevel): +class SparseCOOLevel(AbstractLevel): """Coordinate (COO) sparse format level storage. This level stores sparse data using N coordinate lists (one per dimension). @@ -382,7 +382,7 @@ class SparseCOO(AbstractLevel): -------- Create a 2D sparse tensor in COO format: - >>> coo_2d = SparseCOO(2, Element(0.0)) + >>> coo_2d = SparseCOOLevel(2, ElementLevel(0.0)) """ def __init__(self, ndim, lvl, dims=None, tbl : tuple[Buffer] | None =None): args = [lvl._obj] @@ -401,7 +401,7 @@ def tbl(self) -> tuple[Buffer, ...]: return tuple(jlobj_to_buffer(coord) for coord in self._obj.tbl) -class SparseHash(AbstractLevel): +class SparseHashLevel(AbstractLevel): """Hash table based sparse format level storage. Uses a hash table to store sparse data, supporting efficient random access @@ -422,7 +422,7 @@ class SparseHash(AbstractLevel): -------- Create a 2D sparse tensor using hash storage: - >>> hash_2d = SparseHash(2, Element(0.0)) + >>> hash_2d = SparseHashLevel(2, ElementLevel(0.0)) """ def __init__(self, ndim, lvl, dims=None): args = [lvl._obj] @@ -459,11 +459,11 @@ def construct_levels(obj: JuliaObj, fill_value: number) -> LevelFType: If an unsupported level type is encountered. """ if jl.isa(obj.lvl, jl.Finch.Element): - return Element(fill_value) + return ElementLevel(fill_value) if jl.isa(obj.lvl, jl.Finch.Dense): - return Dense(construct_levels(obj.lvl, fill_value)) + return DenseLevel(construct_levels(obj.lvl, fill_value)) if jl.isa(obj.lvl, jl.Finch.SparseList): - return SparseList(construct_levels(obj.lvl, fill_value)) + return SparseListLevel(construct_levels(obj.lvl, fill_value)) if jl.isa(obj.lvl, jl.Finch.SparseByteMap): - return SparseByteMap(construct_levels(obj.lvl, fill_value)) + return SparseByteMapLevel(construct_levels(obj.lvl, fill_value)) raise Exception("Unhandled exception!") diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py index d9f4d01..2130b43 100644 --- a/src/finch/scheduler.py +++ b/src/finch/scheduler.py @@ -11,7 +11,7 @@ from finchlite.finch_logic import LogicLoader from .compiler import FinchJLCompiler -from .levels import Dense, Element +from .levels import DenseLevel, ElementLevel from .tensor import FinchJLTensorFType @@ -23,9 +23,9 @@ def __init__( super().__init__(loader) def get_output_tns_ftype(self, fill_value: Any, shape_type: tuple[Any, ...]): - lvl = Element(fill_value) + lvl = ElementLevel(fill_value) for _ in shape_type: - lvl = Dense(lvl) + lvl = DenseLevel(lvl) return FinchJLTensorFType(lvl) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index ef2da6c..fdda60c 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -6,7 +6,7 @@ from .julia import jc, jl from .levels import LevelFType, construct_levels -from .typing import JuliaObj +from .typing import JuliaObj, DType from .utils import add_missing_dims, add_plus_one, expand_ellipsis # Tensor Class and associated ftype @@ -160,16 +160,15 @@ def asarray( elif isinstance(obj, np.ndarray): if copy: if np.isfortran(obj): - arr = arr.copy() + obj = obj.copy() else: obj = np.asfortranarray(obj) - dtype = arr.dtype.type - if ( - dtype == np.bool_ - ): # Fails with: Finch currently only supports isbits defaults - dtype = jl_dtypes.bool - lvl = ElementLevel(fill_value, arr.reshape(-1, order="F")) - for i in arr.shape: + else: + if not np.isfortran(obj): + obj = np.asfortranarray(obj) + + lvl = ElementLevel(fill_value, NumpyBuffer(obj.reshape(-1))) + for i in obj.shape: lvl = DenseLevel(lvl, i) return FinchJLTensor(lvl) elif hasattr(x, "__module__") and x.__module__.startswith("scipy.sparse"): @@ -196,12 +195,12 @@ def asarray( ElementLevel( dtype, fill_value, - obj.data + NumpyBUffer(obj.data) ), 2, idxs = ( - x.cols, - x.rows, + PlusOneBuffer(NumpyBuuffer(x.cols)), + PlusOneBuffer(NumpyBuuffer(x.rows)), ), ) ) @@ -215,8 +214,8 @@ def asarray( obj.data ), n, - obj.indptr, - obj.indices + PlusOneBuffer(obj.indptr), + PlusOneBuffer(obj.indices) ), (m, n) ) @@ -225,6 +224,8 @@ def asarray( raise ValueError(f"Unsupported SciPy format: {type(x)}") else: raise ValueError( - "Either scalar, numpy, scipy.sparse or a raw julia object should " + "Either numpy array or a Finch tensor should " f"be provided. Found: {type(obj)}" ) + + diff --git a/tests/conftest.py b/tests/conftest.py index 455e172..51d0cb4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,7 +2,6 @@ import numpy as np - @pytest.fixture def rng(): return np.random.default_rng(42) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 0191f69..510fede 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -29,10 +29,10 @@ from finch.compiler import FinchJLCompiler, FinchJLKernel from finch.julia import jl -from finch.levels import Dense, Element +from finch.levels import DenseLevel, ElementLevel from finch.tensor import FinchJLTensor -a_format = Dense(Dense(Element(0))) +a_format = DenseLevel(DenseLevel(ElementLevel(0))) @pytest.mark.parametrize( diff --git a/tests/test_levels.py b/tests/test_levels.py index 80e8f04..c4df3e5 100644 --- a/tests/test_levels.py +++ b/tests/test_levels.py @@ -3,15 +3,15 @@ import pytest import numpy as np from finch.levels import ( - Element, - Dense, - Pattern, - SparseList, - SparseByteMap, - RepeatRLE, - SparseVBL, - SparseCOO, - SparseHash, + ElementLevel, + DenseLevel, + PatternLevel, + SparseListLevel, + SparseByteMapLevel, + RepeatRLELevel, + SparseVBLLevel, + SparseCOOLevel, + SparseHashLevel, ) from finch.buffer import NumpyBuffer @@ -21,17 +21,17 @@ class TestElement: def test_element_creation_basic(self): """Test creating an Element level with a fill value.""" - elem = Element(0.0) + elem = ElementLevel(0.0) assert elem._obj is not None def test_element_creation_with_int_fill(self): """Test creating an Element level with integer fill value.""" - elem = Element(0) + elem = ElementLevel(0) assert elem._obj is not None def test_element_creation_with_float_fill(self): """Test creating an Element level with float fill value.""" - elem = Element(3.14) + elem = ElementLevel(3.14) assert elem._obj is not None @@ -40,21 +40,21 @@ class TestDense: def test_dense_creation_basic(self): """Test creating a Dense level.""" - elem = Element(0.0) - dense = Dense(elem) + elem = ElementLevel(0.0) + dense = DenseLevel(elem) assert dense._obj is not None def test_dense_creation_with_shape(self): """Test creating a Dense level with explicit shape.""" - elem = Element(0.0) - dense = Dense(elem, shape=10) + elem = ElementLevel(0.0) + dense = DenseLevel(elem, shape=10) assert dense._obj is not None def test_dense_nesting(self): """Test creating nested Dense levels.""" - elem = Element(0.0) - dense1 = Dense(elem) - dense2 = Dense(dense1) + elem = ElementLevel(0.0) + dense1 = DenseLevel(elem) + dense2 = DenseLevel(dense1) assert dense2._obj is not None @@ -63,7 +63,7 @@ class TestPattern: def test_pattern_creation(self): """Test creating a Pattern level.""" - pattern = Pattern() + pattern = PatternLevel() assert pattern._obj is not None @@ -72,44 +72,44 @@ class TestSparseList: def test_sparselist_creation_basic(self): """Test creating a SparseList level.""" - elem = Element(0.0) - sparse = SparseList(elem) + elem = ElementLevel(0.0) + sparse = SparseListLevel(elem) assert sparse._obj is not None def test_sparselist_creation_with_dim(self): """Test creating a SparseList level with explicit dimension.""" - elem = Element(0.0) - sparse = SparseList(elem, dim=10) + elem = ElementLevel(0.0) + sparse = SparseListLevel(elem, dim=10) assert sparse._obj is not None def test_sparselist_creation_with_data_arrays(self): """Test creating a SparseList level with pointer and index arrays.""" - elem = Element(0.0) + elem = ElementLevel(0.0) ptr = NumpyBuffer(np.array([0, 2, 2, 4], dtype=np.int32)) idx = NumpyBuffer(np.array([1, 2, 1, 3], dtype=np.int32)) - sparse = SparseList(elem, ptr=ptr, idx=idx) + sparse = SparseListLevel(elem, ptr=ptr, idx=idx) assert sparse._obj is not None def test_sparselist_creation_with_data_lists(self): """Test creating a SparseList level with pointer and index as lists.""" - elem = Element(0.0) + elem = ElementLevel(0.0) ptr = [0, 2, 2, 4] idx = [1, 2, 1, 3] - sparse = SparseList(elem, ptr=ptr, idx=idx) + sparse = SparseListLevel(elem, ptr=ptr, idx=idx) assert sparse._obj is not None def test_sparselist_ptr_property(self): """Test accessing ptr property of SparseList.""" - elem = Element(0.0) - sparse = SparseList(elem) + elem = ElementLevel(0.0) + sparse = SparseListLevel(elem) # Property should be accessible ptr_buffer = sparse.ptr assert ptr_buffer is not None def test_sparselist_idx_property(self): """Test accessing idx property of SparseList.""" - elem = Element(0.0) - sparse = SparseList(elem) + elem = ElementLevel(0.0) + sparse = SparseListLevel(elem) # Property should be accessible idx_buffer = sparse.idx assert idx_buffer is not None @@ -120,14 +120,14 @@ class TestSparseByteMap: def test_sparsebytemap_creation_basic(self): """Test creating a SparseByteMap level.""" - elem = Element(0.0) - sparse = SparseByteMap(elem) + elem = ElementLevel(0.0) + sparse = SparseByteMapLevel(elem) assert sparse._obj is not None def test_sparsebytemap_creation_with_dim(self): """Test creating a SparseByteMap level with explicit dimension.""" - elem = Element(0.0) - sparse = SparseByteMap(elem, dim=10) + elem = ElementLevel(0.0) + sparse = SparseByteMapLevel(elem, dim=10) assert sparse._obj is not None @@ -136,14 +136,14 @@ class TestRepeatRLE: def test_repeatrle_creation_basic(self): """Test creating a RepeatRLE level.""" - elem = Element(0.0) - rle = RepeatRLE(elem) + elem = ElementLevel(0.0) + rle = RepeatRLELevel(elem) assert rle._obj is not None def test_repeatrle_creation_with_dim(self): """Test creating a RepeatRLE level with explicit dimension.""" - elem = Element(0.0) - rle = RepeatRLE(elem, dim=10) + elem = ElementLevel(0.0) + rle = RepeatRLELevel(elem, dim=10) assert rle._obj is not None @@ -152,14 +152,14 @@ class TestSparseVBL: def test_sparsevbl_creation_basic(self): """Test creating a SparseVBL level.""" - elem = Element(0.0) - vbl = SparseVBL(elem) + elem = ElementLevel(0.0) + vbl = SparseVBLLevel(elem) assert vbl._obj is not None def test_sparsevbl_creation_with_dim(self): """Test creating a SparseVBL level with explicit dimension.""" - elem = Element(0.0) - vbl = SparseVBL(elem, dim=10) + elem = ElementLevel(0.0) + vbl = SparseVBLLevel(elem, dim=10) assert vbl._obj is not None @@ -168,48 +168,48 @@ class TestSparseCOO: def test_sparsecoo_creation_basic(self): """Test creating a SparseCOO level.""" - elem = Element(0.0) - coo = SparseCOO(2, elem) + elem = ElementLevel(0.0) + coo = SparseCOOLevel(2, elem) assert coo._obj is not None def test_sparsecoo_creation_with_dims(self): """Test creating a SparseCOO level with explicit dimensions.""" - elem = Element(0.0) - coo = SparseCOO(2, elem, dims=(4, 3)) + elem = ElementLevel(0.0) + coo = SparseCOOLevel(2, elem, dims=(4, 3)) assert coo._obj is not None def test_sparsecoo_creation_with_dims_list(self): """Test creating a SparseCOO level with dimensions as list.""" - elem = Element(0.0) - coo = SparseCOO(2, elem, dims=[4, 3]) + elem = ElementLevel(0.0) + coo = SparseCOOLevel(2, elem, dims=[4, 3]) assert coo._obj is not None def test_sparsecoo_creation_with_coordinate_arrays(self): """Test creating a SparseCOO level with coordinate arrays.""" - elem = Element(0.0) + elem = ElementLevel(0.0) i_coords = NumpyBuffer(np.array([0, 1, 2, 3], dtype=np.int32)) j_coords = NumpyBuffer(np.array([0, 0, 2, 2], dtype=np.int32)) - coo = SparseCOO(2, elem, tbl=(i_coords, j_coords)) + coo = SparseCOOLevel(2, elem, tbl=(i_coords, j_coords)) assert coo._obj is not None def test_sparsecoo_creation_with_coordinate_lists(self): """Test creating a SparseCOO level with coordinate arrays as lists.""" - elem = Element(0.0) + elem = ElementLevel(0.0) i_coords = [0, 1, 2, 3] j_coords = [0, 0, 2, 2] - coo = SparseCOO(2, elem, tbl=(i_coords, j_coords)) + coo = SparseCOOLevel(2, elem, tbl=(i_coords, j_coords)) assert coo._obj is not None def test_sparsecoo_3d(self): """Test creating a 3D SparseCOO level.""" - elem = Element(0.0) - coo = SparseCOO(3, elem, dims=(5, 4, 3)) + elem = ElementLevel(0.0) + coo = SparseCOOLevel(3, elem, dims=(5, 4, 3)) assert coo._obj is not None def test_sparsecoo_tbl_property(self): """Test accessing tbl property of SparseCOO.""" - elem = Element(0.0) - coo = SparseCOO(2, elem) + elem = ElementLevel(0.0) + coo = SparseCOOLevel(2, elem) # Property should be accessible tbl = coo.tbl assert tbl is not None @@ -221,26 +221,26 @@ class TestSparseHash: def test_sparsehash_creation_basic(self): """Test creating a SparseHash level.""" - elem = Element(0.0) - hash_level = SparseHash(2, elem) + elem = ElementLevel(0.0) + hash_level = SparseHashLevel(2, elem) assert hash_level._obj is not None def test_sparsehash_creation_with_dims(self): """Test creating a SparseHash level with explicit dimensions.""" - elem = Element(0.0) - hash_level = SparseHash(2, elem, dims=(4, 3)) + elem = ElementLevel(0.0) + hash_level = SparseHashLevel(2, elem, dims=(4, 3)) assert hash_level._obj is not None def test_sparsehash_creation_with_dims_list(self): """Test creating a SparseHash level with dimensions as list.""" - elem = Element(0.0) - hash_level = SparseHash(2, elem, dims=[4, 3]) + elem = ElementLevel(0.0) + hash_level = SparseHashLevel(2, elem, dims=[4, 3]) assert hash_level._obj is not None def test_sparsehash_3d(self): """Test creating a 3D SparseHash level.""" - elem = Element(0.0) - hash_level = SparseHash(3, elem, dims=(5, 4, 3)) + elem = ElementLevel(0.0) + hash_level = SparseHashLevel(3, elem, dims=(5, 4, 3)) assert hash_level._obj is not None @@ -249,32 +249,32 @@ class TestComposedLevels: def test_csc_matrix_format(self): """Test creating CSC matrix format (Dense(SparseList(Element))).""" - elem = Element(0.0) - sparse = SparseList(elem) - dense = Dense(sparse) + elem = ElementLevel(0.0) + sparse = SparseListLevel(elem) + dense = DenseLevel(sparse) assert dense._obj is not None def test_csr_like_format(self): """Test creating CSR-like format (SparseList(Dense(Element))).""" - elem = Element(0.0) - dense = Dense(elem) - sparse = SparseList(dense) + elem = ElementLevel(0.0) + dense = DenseLevel(elem) + sparse = SparseListLevel(dense) assert sparse._obj is not None def test_dcsc_format(self): """Test creating DCSC format (SparseList(SparseList(Element))).""" - elem = Element(0.0) - sparse1 = SparseList(elem) - sparse2 = SparseList(sparse1) + elem = ElementLevel(0.0) + sparse1 = SparseListLevel(elem) + sparse2 = SparseListLevel(sparse1) assert sparse2._obj is not None def test_deep_nesting(self): """Test deeply nested levels.""" - elem = Element(0.0) - sparse = SparseList(elem) - dense = Dense(sparse) - sparse2 = SparseList(dense) - dense2 = Dense(sparse2) + elem = ElementLevel(0.0) + sparse = SparseListLevel(elem) + dense = DenseLevel(sparse) + sparse2 = SparseListLevel(dense) + dense2 = DenseLevel(sparse2) assert dense2._obj is not None @@ -283,34 +283,34 @@ class TestEdgeCases: def test_element_with_negative_fill(self): """Test Element with negative fill value.""" - elem = Element(-1.0) + elem = ElementLevel(-1.0) assert elem._obj is not None def test_sparselist_only_ptr_no_idx(self): """Test SparseList with ptr but no idx (should not add both).""" - elem = Element(0.0) + elem = ElementLevel(0.0) ptr = NumpyBuffer(np.array([0, 2, 2], dtype=np.int32)) - sparse = SparseList(elem, ptr=ptr) + sparse = SparseListLevel(elem, ptr=ptr) assert sparse._obj is not None def test_sparselist_only_idx_no_ptr(self): """Test SparseList with idx but no ptr (should not add both).""" - elem = Element(0.0) + elem = ElementLevel(0.0) idx = NumpyBuffer(np.array([1, 2], dtype=np.int32)) - sparse = SparseList(elem, idx=idx) + sparse = SparseListLevel(elem, idx=idx) assert sparse._obj is not None def test_sparsecoo_single_coordinate(self): """Test SparseCOO with single coordinate.""" - elem = Element(0.0) + elem = ElementLevel(0.0) coords = NumpyBuffer(np.array([0], dtype=np.int32)) - coo = SparseCOO(1, elem, tbl=(coords,)) + coo = SparseCOOLevel(1, elem, tbl=(coords,)) assert coo._obj is not None def test_large_dimension(self): """Test levels with large dimensions.""" - elem = Element(0.0) - sparse = SparseList(elem, dim=1000000) + elem = ElementLevel(0.0) + sparse = SparseListLevel(elem, dim=1000000) assert sparse._obj is not None @@ -319,28 +319,28 @@ class TestArrayConversion: def test_sparselist_converts_lists_to_arrays(self): """Test that SparseList converts list arguments to arrays.""" - elem = Element(0.0) + elem = ElementLevel(0.0) ptr = NumpyBuffer(np.array([0, 2, 4], dtype=np.int32)) idx = NumpyBuffer(np.array([1, 2, 3], dtype=np.int32)) - sparse = SparseList(elem, ptr=ptr, idx=idx) + sparse = SparseListLevel(elem, ptr=ptr, idx=idx) # Should not raise an error during creation assert sparse._obj is not None def test_sparsecoo_converts_lists_to_arrays(self): """Test that SparseCOO converts list arguments to arrays.""" - elem = Element(0.0) + elem = ElementLevel(0.0) coords_list = ( NumpyBuffer(np.array([0, 1, 2], dtype=np.int32)), NumpyBuffer(np.array([0, 1, 2], dtype=np.int32)) ) - coo = SparseCOO(2, elem, tbl=coords_list) + coo = SparseCOOLevel(2, elem, tbl=coords_list) # Should not raise an error during creation assert coo._obj is not None def test_different_dtype_arrays(self): """Test that different dtype arrays are handled.""" - elem = Element(0.0) + elem = ElementLevel(0.0) ptr = NumpyBuffer(np.array([0, 2, 4], dtype=np.int64)) idx = NumpyBuffer(np.array([1, 2, 3], dtype=np.int32)) - sparse = SparseList(elem, ptr=ptr, idx=idx) + sparse = SparseListLevel(elem, ptr=ptr, idx=idx) assert sparse._obj is not None From 790816addea8a770ebf8734871fcd6f9232ae218 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 19:52:36 -0400 Subject: [PATCH 46/81] fix --- src/finch/levels.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 705526d..493a6c0 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -149,7 +149,17 @@ def __str__(self): class AbstractLevel(_Display): - pass + @property + def ndim(self) -> int: + return self.ftype().ndim + + @property + def fill_value(self) -> Any: + return self.ftype().fill_value + + @property + def element_type(self) -> Any: + return self.ftype().element_type # core levels From 33b8eb211fbc45fc3e9f1f05f6f8414d2689d350 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 19:56:27 -0400 Subject: [PATCH 47/81] fix --- pytest.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytest.ini b/pytest.ini index e00b2ce..03faf57 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,5 +1,5 @@ [pytest] -addopts = --cov-report term-missing --cov-report html --cov-report=xml --cov-report=term --cov finch --cov-config .coveragerc --junitxml=junit/test-results.xml +addopts = --capture=tee-sys --cov-report term-missing --cov-report html --cov-report=xml --cov-report=term --cov finch --cov-config .coveragerc --junitxml=junit/test-results.xml filterwarnings = ignore::PendingDeprecationWarning testpaths = From 7623e84d0374cac74620a0b49e9c8ddc5f4c416c Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 20:01:40 -0400 Subject: [PATCH 48/81] fix --- src/finch/levels.py | 2 +- src/finch/tensor.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 493a6c0..82f68cd 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -445,7 +445,7 @@ def __init__(self, ndim, lvl, dims=None): # Helper Methods -def construct_levels(obj: JuliaObj, fill_value: number) -> LevelFType: +def jlobj_to_level(obj: JuliaObj, fill_value: number) -> LevelFType: """Construct a level hierarchy from a Julia Finch tensor object. Recursively constructs a Python representation of the tensor's level structure diff --git a/src/finch/tensor.py b/src/finch/tensor.py index fdda60c..ec48dac 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -5,7 +5,7 @@ from finchlite import EagerTensor, Tensor, TensorFType from .julia import jc, jl -from .levels import LevelFType, construct_levels +from .levels import LevelFType, ElementLevel, DenseLevel, SparseListLevel, SparseCOOLevel from .typing import JuliaObj, DType from .utils import add_missing_dims, add_plus_one, expand_ellipsis From bd09763982c6f50b9fd1d34afbcdf086f8fbcddc Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 20:04:26 -0400 Subject: [PATCH 49/81] fix --- src/finch/tensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index ec48dac..5b1efbf 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -171,7 +171,7 @@ def asarray( for i in obj.shape: lvl = DenseLevel(lvl, i) return FinchJLTensor(lvl) - elif hasattr(x, "__module__") and x.__module__.startswith("scipy.sparse"): + elif hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): if obj.format == "coo": obj = obj.T if copy: From e881113566839540fb29e59cb70f8f450f60939b Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 26 Mar 2026 23:29:11 -0400 Subject: [PATCH 50/81] fix --- src/finch/levels.py | 390 ++++++-------------------------------------- src/finch/tensor.py | 40 ++--- 2 files changed, 66 insertions(+), 364 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 82f68cd..6047be5 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -3,8 +3,7 @@ import numpy as np -from finchlite import Tensor, TensorFType -from finchlite.finch_assembly import Buffer, BufferFType +from finchlite.finch_assembly import Buffer, BufferFormat from .buffer import PlusOneBuffer, NumpyBuffer, buffer_to_jlobj, jlobj_to_buffer from .julia import jl @@ -12,21 +11,14 @@ from dataclasses import dataclass -# Abstract FTypes -class LevelFType(TensorFType): - def from_numpy(self, _) -> Tensor: - raise NotImplementedError - +# Abstract Formats +class LevelFormat(): @property - def shape_type(self) -> tuple[type, ...]: - return tuple(self.element_type for _ in range(self.ndim)) - - def __call__(self, _) -> Tensor: - raise Exception("Cannot create an object of this type!") - + @abstractmethod + def shape_type(self) -> tuple: ... -class NestedLevelFType(LevelFType): +class NestedLevelFormat(LevelFormat): @property def ndim(self) -> np.intp: return self.lvl.ndim + np.intp(1) @@ -49,7 +41,7 @@ def __hash__(self): def create_jl_obj(self) -> JuliaObj: ... -class ElementFType(LevelFType): +class ElementFormat(LevelFormat): """Element level storage format for scalar tensor leaves. A subfiber of an element level is a scalar, initialized to a fill value. @@ -71,7 +63,7 @@ def element_type(self) -> Any: return type(self._fill_value) def __eq__(self, other): - return isinstance(other, ElementFType) and self._fill_value == other.fill_value + return isinstance(other, ElementFormat) and self._fill_value == other.fill_value def __hash__(self): return hash((self.__class__.__name__, self._fill_value)) @@ -79,36 +71,45 @@ def __hash__(self): def create_jl_obj(self) -> JuliaObj: return jl.Element(self._fill_value) - + def shape_type(self) -> tuple: + return () @dataclass(frozen=True) -class DenseFType(NestedLevelFType): +class DenseFormat(NestedLevelFormat): """Dense format wrapper type for Finch tensors. A dense format stores every slice of a tensor. A subfiber of a dense level is an array which stores every slice A[:, ..., :, i] as a distinct subfiber. Dense levels support both row-major and column-major access. """ - lvl: NestedLevelFType + lvl: NestedLevelFormat + dim_type: type = np.intp def create_jl_obj(self) -> JuliaObj: return jl.Dense(self.lvl.create_jl_obj()) + + def shape_type(self) -> tuple: + return self.lvl.shape_type + (self.dim_type,) @dataclass(frozen=True) -class SparseListFType(NestedLevelFType): +class SparseListFormat(NestedLevelFormat): """Sparse list format wrapper type for Finch tensors. A sparse list format stores only potentially non-fill slices using a sorted list. Slices that are entirely fill_value are omitted. This format is efficient for tensors with sparse patterns and supports column-major reads and bulk updates. """ - lvl: NestedLevelFType + lvl: NestedLevelFormat + dim_type: type = np.intp def create_jl_obj(self) -> JuliaObj: return jl.SparseList(self.lvl.create_jl_obj()) + def shape_type(self) -> tuple: + return self.lvl.shape_type + (self.dim_type,) + @dataclass(frozen=True) -class SparseCOOFType(NestedLevelFType): +class SparseCOOFormat(NestedLevelFormat): """Coordinate (COO) format wrapper type for Finch tensors. A coordinate format stores sparse tensors as lists of coordinates. @@ -116,336 +117,37 @@ class SparseCOOFType(NestedLevelFType): with coordinates sorted in column-major order. This is a legacy format maintained for backward compatibility. """ - lvl: NestedLevelFType + lvl: NestedLevelFormat N: int = 2 + dim_type: Tuple{type} | None = np.intp + def create_jl_obj(self) -> JuliaObj: return jl.SparseCOO(self.lvl.create_jl_obj()) + def shape_type(self) -> tuple: + if self.dim_type is None: + return self.lvl.shape_type + (self.N * self.lvl.ndim,) + return self.lvl.shape_type + self.dim_type + @dataclass(frozen=True) -class SparseByteMapFType(NestedLevelFType): +class SparseByteMapFormat(NestedLevelFormat): """Sparse byte map format wrapper type for Finch tensors. A sparse byte map format uses a dense bitmap to encode which slices are stored, similar to SparseList but supporting random access. Only potentially non-fill slices are stored as subfibers. """ - lvl: NestedLevelFType + lvl: NestedLevelFormat + dim_type: type = np.intp def create_jl_obj(self) -> JuliaObj: return jl.SparseByteMap(self.lvl.create_jl_obj()) - - -class _Display: - _obj: JuliaObj - - def __repr__(self): - return jl.sprint(jl.show, self._obj) - - def __str__(self): - return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) - - -# LEVEL - - -class AbstractLevel(_Display): - @property - def ndim(self) -> int: - return self.ftype().ndim - - @property - def fill_value(self) -> Any: - return self.ftype().fill_value - - @property - def element_type(self) -> Any: - return self.ftype().element_type - - -# core levels - - -class DenseLevel(AbstractLevel): - """Dense level storage format. - - A subfiber of a dense level is an array which stores every slice as a distinct - subfiber in the child level. Dense levels support efficient random access and - both column-major and out-of-order updates. - - Parameters - ---------- - lvl : AbstractLevel - The child level that will store each slice. - shape : int, optional - The size of the dimension at this level. If not provided, the shape - is inferred from the child level. - - Examples - -------- - Create a 2D dense tensor: - - >>> dense_2d = DenseLevel(DenseLevel(ElementLevel(0.0))) - """ - def __init__(self, lvl, shape=None): - args = [lvl._obj] - if shape is not None: - args.append(shape) - self._obj = jl.Dense(*args) - - -class ElementLevel(AbstractLevel): - """Element level storage format (leaf level). - - A subfiber of an element level is a scalar of a specified type, initialized - to a fill value. Element levels form the leaf nodes of the tensor tree and - store the actual data values. - - Parameters - ---------- - fill_value : float or int - The default fill value for elements that are not explicitly set. - data : array-like, optional - Optional vector to store the element data. If not provided, a new - vector is created. - - Examples - -------- - Create an element level with zero fill value: - - >>> elem = ElementLevel(0.0) - """ - def __init__(self, fill_value, data: Buffer | None=None): - args = [fill_value] - if data is not None: - args.append(buffer_to_jlobj(data)) - self._obj = jl.Element(*args) - - @property - def data(self) -> Buffer: - """Return the data buffer for this element level.""" - return jlobj_to_buffer(self._obj.data) - - -class PatternLevel(AbstractLevel): - """Pattern level storage format (leaf level). - - A subfiber of a pattern level is the boolean value true, but with a fill - value of false. Pattern levels are used to create tensors representing - which values are stored by other fibers, allowing for tracking sparsity - structure without storing values. - - Examples - -------- - Create a pattern level to track which elements are nonzero: - - >>> pattern = PatternLevel() - """ - def __init__(self): - self._obj = jl.Pattern() - - -# advanced levels - - -class SparseListLevel(AbstractLevel): - """Sparse list level storage format. - - A subfiber of a sparse list level stores only potentially non-fill slices - using a sorted list to record which slices are stored. Slices that are - entirely fill_value are omitted. This format is efficient for sparse tensors - and supports column-major reads and bulk updates, but does not support - random access or random updates. - - Parameters - ---------- - lvl : AbstractLevel - The child level that will store each non-fill slice. - dim : int, optional - The size of the dimension at this level. If not provided, the dimension - is inferred from the child level. - ptr : array-like, optional - Array of positions/pointers for the sparse list. If not provided, - will be created internally. Converted to ndarray if provided. - idx : array-like, optional - Array of indices for the sparse list. If not provided, - will be created internally. Converted to ndarray if provided. - - Examples - -------- - Create a sparse matrix in CSC format: - - >>> sparse_matrix = SparseListLevel(SparseListLevel(ElementLevel(0.0))) - """ - def __init__(self, lvl, dim=None, ptr=None, idx=None): - args = [lvl._obj] - if dim is not None: - args.append(dim) - if ptr is not None and idx is not None: - args.append(buffer_to_jlobj(ptr)) - args.append(buffer_to_jlobj(idx)) - self._obj = jl.SparseList(*args) - - @property - def ptr(self) -> Buffer: - """Return the pointer array for this sparse list level.""" - return jlobj_to_buffer(self._obj.ptr) - - @property - def idx(self) -> Buffer: - """Return the index array for this sparse list level.""" - return jlobj_to_buffer(self._obj.idx) - -class SparseByteMapLevel(AbstractLevel): - """Sparse byte map level storage format. - Similar to SparseList, but uses a dense bitmap to encode which slices are - stored instead of a sparse index list. This allows for efficient random - access while still omitting fill_value slices. The byte map approach trades - memory for faster lookups. - - Parameters - ---------- - lvl : AbstractLevel - The child level that will store each non-fill slice. - dim : int, optional - The size of the dimension at this level. If not provided, the dimension - is inferred from the child level. - - Examples - -------- - Create a sparse matrix with byte map storage: - - >>> sparse_matrix = SparseByteMapLevel(SparseByteMapLevel(ElementLevel(0.0))) - """ - def __init__(self, lvl, dim=None): - args = [lvl._obj] - if dim is not None: - args.append(dim) - self._obj = jl.SparseByteMap(*args) - - -class RepeatRLELevel(AbstractLevel): - """Run-length encoding level for repeated values. - - This level stores runs of repeated values, useful for data with - long sequences of identical values. - - Parameters - ---------- - lvl : AbstractLevel - The child level to store run information. - dim : int, optional - The size of the dimension at this level. If not provided, the dimension - is inferred from the child level. - """ - def __init__(self, lvl, dim=None): - args = [lvl._obj] - if dim is not None: - args.append(dim) - self._obj = jl.RepeatRLE(*args) - - -class SparseVBLLevel(AbstractLevel): - """Sparse variable block list level storage format. - - Like SparseList, but stores contiguous slices together in blocks. - This can improve cache locality and performance for certain access patterns. - - Parameters - ---------- - lvl : AbstractLevel - The child level that will store each block of slices. - dim : int, optional - The size of the dimension at this level. If not provided, the dimension - is inferred from the child level. - """ - def __init__(self, lvl, dim=None): - args = [lvl._obj] - if dim is not None: - args.append(dim) - self._obj = jl.SparseVBL(*args) - - -class SparseCOOLevel(AbstractLevel): - """Coordinate (COO) sparse format level storage. - - This level stores sparse data using N coordinate lists (one per dimension). - Coordinates are stored in column-major order. This is a legacy format - maintained for backward compatibility; consider using other sparse formats - for new code. - - Parameters - ---------- - ndim : int - Number of dimensions for the coordinate format. - lvl : AbstractLevel - The child level to store coordinate data. - dims : tuple of int, optional - Sizes of the last N dimensions. If not provided, dimensions are - inferred from the child level. - tbl : tuple of arrays, optional - A tuple of coordinate arrays (one per dimension). Each array stores - the coordinates for that dimension. Arrays are converted to ndarray - if provided. If not provided, will be created internally. - - Examples - -------- - Create a 2D sparse tensor in COO format: - - >>> coo_2d = SparseCOOLevel(2, ElementLevel(0.0)) - """ - def __init__(self, ndim, lvl, dims=None, tbl : tuple[Buffer] | None =None): - args = [lvl._obj] - if dims is not None: - if isinstance(dims, (list, tuple)): - args.extend(dims) - else: - args.append(dims) - if tbl is not None: - args.extend([buffer_to_jlobj(t) for t in tbl]) - self._obj = jl.SparseCOO[ndim](*args) - - @property - def tbl(self) -> tuple[Buffer, ...]: - """Return the coordinate array tuple for this COO level.""" - return tuple(jlobj_to_buffer(coord) for coord in self._obj.tbl) - - -class SparseHashLevel(AbstractLevel): - """Hash table based sparse format level storage. - - Uses a hash table to store sparse data, supporting efficient random access - and random updates. This format is useful when you need flexible out-of-order - insertion of elements. - - Parameters - ---------- - ndim : int - Number of dimensions for the hash format. - lvl : AbstractLevel - The child level to store hash table data. - dims : tuple of int, optional - Sizes of the last N dimensions. If not provided, dimensions are - inferred from the child level. - - Examples - -------- - Create a 2D sparse tensor using hash storage: - - >>> hash_2d = SparseHashLevel(2, ElementLevel(0.0)) - """ - def __init__(self, ndim, lvl, dims=None): - args = [lvl._obj] - if dims is not None: - if isinstance(dims, (list, tuple)): - args.extend(dims) - else: - args.append(dims) - self._obj = jl.SparseHash[ndim](*args) - + def shape_type(self) -> tuple: + return self.lvl.shape_type + (self.dim_type,) # Helper Methods -def jlobj_to_level(obj: JuliaObj, fill_value: number) -> LevelFType: +def jlobj_to_format(obj: JuliaObj) -> LevelFType: """Construct a level hierarchy from a Julia Finch tensor object. Recursively constructs a Python representation of the tensor's level structure @@ -468,12 +170,12 @@ def jlobj_to_level(obj: JuliaObj, fill_value: number) -> LevelFType: Exception If an unsupported level type is encountered. """ - if jl.isa(obj.lvl, jl.Finch.Element): - return ElementLevel(fill_value) - if jl.isa(obj.lvl, jl.Finch.Dense): - return DenseLevel(construct_levels(obj.lvl, fill_value)) - if jl.isa(obj.lvl, jl.Finch.SparseList): - return SparseListLevel(construct_levels(obj.lvl, fill_value)) - if jl.isa(obj.lvl, jl.Finch.SparseByteMap): - return SparseByteMapLevel(construct_levels(obj.lvl, fill_value)) - raise Exception("Unhandled exception!") + if jl.isa(obj, jl.Finch.Element): + return ElementLevel(jl.fill_value(obj)) + if jl.isa(obj, jl.Finch.Dense): + return DenseLevel(type(obj.shape), jlobj_to_format(obj.lvl)) + if jl.isa(obj, jl.Finch.SparseList): + return SparseListLevel(type(obj.shape), jlobj_to_format(obj.lvl)) + if jl.isa(obj, jl.Finch.SparseByteMap): + return SparseByteMapLevel(jlobj_to_format(obj.lvl)) + raise Exception("Unhandled exception!") \ No newline at end of file diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 5b1efbf..23d5754 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -5,14 +5,14 @@ from finchlite import EagerTensor, Tensor, TensorFType from .julia import jc, jl -from .levels import LevelFType, ElementLevel, DenseLevel, SparseListLevel, SparseCOOLevel +from .levels import LevelFType, ElementLevel, DenseLevel, SparseListLevel, SparseCOOLevel, jlobj_to_format from .typing import JuliaObj, DType from .utils import add_missing_dims, add_plus_one, expand_ellipsis # Tensor Class and associated ftype class FinchJLTensorFType(TensorFType): def __init__(self, lvl): - self._lvl: LevelFType = lvl + self._lvl: LevelFormat = lvl @property def ndim(self) -> np.intp: @@ -49,7 +49,6 @@ def __eq__(self, other): def __hash__(self): return hash(("FinchJLTensorFType", self._lvl)) - class FinchJLTensor(EagerTensor): def __init__(self, obj: JuliaObj): if isinstance(obj, JuliaObj): @@ -62,7 +61,7 @@ def ftype(self) -> TensorFType: """Returns the ftype of the buffer""" if self._is_scalar(): return FinchJLTensorFType(Scalar(self._obj.val)) - return FinchJLTensorFType(construct_levels(self._obj, jl.fill_value(self._obj))) + return FinchJLTensorFType(jlobj_to_format(self._obj, jl.fill_value(self._obj))) @property def shape(self) -> tuple: @@ -169,18 +168,18 @@ def asarray( lvl = ElementLevel(fill_value, NumpyBuffer(obj.reshape(-1))) for i in obj.shape: - lvl = DenseLevel(lvl, i) + lvl = jl.DenseLevel(lvl, i) return FinchJLTensor(lvl) elif hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): if obj.format == "coo": obj = obj.T if copy: if obj.format in ("coo", "csc"): - if not x.has_sorted_indices: + if not obj.has_sorted_indices: obj = obj.sorted_indices() else: obj = obj.copy() - if not x.has_canonical_format: + if not obj.has_canonical_format: obj.sum_duplicates() else: obj = obj.asformat("csc") @@ -191,37 +190,38 @@ def asarray( m, n = obj.shape if obj.format == "coo": return Tensor( - SparseCOOLevel( - ElementLevel( + jl.SparseCOOLevel( + jl.ElementLevel( dtype, fill_value, - NumpyBUffer(obj.data) + obj.data ), 2, idxs = ( - PlusOneBuffer(NumpyBuuffer(x.cols)), - PlusOneBuffer(NumpyBuuffer(x.rows)), + jl.Finch.PlusOneVector(obj.cols), + jl.Finch.PlusOneVector(obj.rows), ), + (m, n) ) ) - elif x.format == "csc": + elif obj.format == "csc": return Tensor( - DenseLevel( - SparseListLevel( - ElementLevel( + jl.DenseLevel( + jl.SparseListLevel( + jl.ElementLevel( dtype, fill_value, obj.data ), n, - PlusOneBuffer(obj.indptr), - PlusOneBuffer(obj.indices) + jl.Finch.PlusOneVector(obj.indptr), + jl.Finch.PlusOneVector(obj.indices) ), - (m, n) + m ) ) else: - raise ValueError(f"Unsupported SciPy format: {type(x)}") + raise ValueError(f"Unsupported SciPy format: {type(obj)}") else: raise ValueError( "Either numpy array or a Finch tensor should " From 9d9ce161ea6e41191d50e3e676149eb59feaef3d Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 17:14:53 -0400 Subject: [PATCH 51/81] fix --- src/finch/__init__.py | 2 +- src/finch/buffer.py | 33 ++-- src/finch/julia.py | 3 +- src/finch/levels.py | 55 ++++--- src/finch/scalar.py | 25 --- src/finch/scheduler.py | 6 +- src/finch/tensor.py | 87 ++++------- tests/conftest.py | 1 + tests/test_asarray.py | 7 +- tests/test_compiler.py | 6 +- tests/test_levels.py | 346 ----------------------------------------- 11 files changed, 99 insertions(+), 472 deletions(-) delete mode 100644 src/finch/scalar.py delete mode 100644 tests/test_levels.py diff --git a/src/finch/__init__.py b/src/finch/__init__.py index 643a54d..f7f916a 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -122,9 +122,9 @@ ) from .scheduler import COMPILE_JULIA from .tensor import ( - asarray, FinchJLTensor, FinchJLTensorFType, + asarray, ) __all__ = [ diff --git a/src/finch/buffer.py b/src/finch/buffer.py index 664cffe..c98f3b0 100644 --- a/src/finch/buffer.py +++ b/src/finch/buffer.py @@ -1,8 +1,12 @@ - from abc import ABC -from finchlite.finch_assembly import Buffer, BufferFType + +import numpy as np + from finchlite.codegen import NumpyBuffer -from .julia import jc, jl +from finchlite.finch_assembly import Buffer, BufferFType + +from .julia import jl + class PlusOneBufferFType(BufferFType): def __init__(self, data_ftype): @@ -10,11 +14,11 @@ def __init__(self, data_ftype): def __call__(self, *args, **kwargs): return PlusOneBuffer(self.data_ftype(*args, **kwargs)) - + @property def element_type(self): return self.data_ftype.element_type - + @property def length_type(self): return self.data_ftype.length_type @@ -22,12 +26,13 @@ def length_type(self): class PlusOneBuffer(Buffer, ABC): """ - Buffer that adds one to each element when loaded and subtracts one from each element when stored. + Buffer that adds one to each element when loaded and subtracts one from each + element when stored. """ def __init__(self, data): - self.data : Buffer = data - + self.data: Buffer = data + def ftype(self): return PlusOneBufferFType(self.data.ftype()) @@ -55,18 +60,18 @@ def store(self, idx: int, val): def resize(self, len: int): self.data.resize(len) + def buffer_to_jlobj(buffer: Buffer): if isinstance(buffer, PlusOneBuffer): return jl.PlusOneVector(buffer_to_jlobj(buffer.data)) - elif isinstance(buffer, NumpyBuffer): + if isinstance(buffer, NumpyBuffer): return buffer.arr - else: - raise ValueError(f"Unsupported buffer type: {type(buffer)}") + raise ValueError(f"Unsupported buffer type: {type(buffer)}") + def jlobj_to_buffer(jlobj): if isinstance(jlobj, jl.PlusOneVector): return PlusOneBuffer(jlobj_to_buffer(jlobj.data)) - elif isinstance(jlobj, np.ndarray): + if isinstance(jlobj, np.ndarray): return NumpyBuffer(jlobj) - else: - raise ValueError(f"Unsupported Julia object type: {type(jlobj)}") \ No newline at end of file + raise ValueError(f"Unsupported Julia object type: {type(jlobj)}") diff --git a/src/finch/julia.py b/src/finch/julia.py index f5f087a..a0a21b3 100644 --- a/src/finch/julia.py +++ b/src/finch/julia.py @@ -1,4 +1,5 @@ -import os # noqa: I001, F401 +import os # noqa: I001, F401 + os.environ["PYTHON_JULIACALL_HANDLE_SIGNALS"] = "yes" import juliapkg # noqa: I001, F401 diff --git a/src/finch/levels.py b/src/finch/levels.py index 6047be5..38063cc 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -1,18 +1,15 @@ from abc import abstractmethod +from dataclasses import dataclass from typing import Any import numpy as np -from finchlite.finch_assembly import Buffer, BufferFormat -from .buffer import PlusOneBuffer, NumpyBuffer, buffer_to_jlobj, jlobj_to_buffer - from .julia import jl from .typing import JuliaObj, number -from dataclasses import dataclass # Abstract Formats -class LevelFormat(): +class LevelFormat: @property @abstractmethod def shape_type(self) -> tuple: ... @@ -43,10 +40,11 @@ def create_jl_obj(self) -> JuliaObj: ... class ElementFormat(LevelFormat): """Element level storage format for scalar tensor leaves. - + A subfiber of an element level is a scalar, initialized to a fill value. The element level is a leaf level used at the end of the tensor tree structure. """ + def __init__(self, fill_value: number): self._fill_value = fill_value @@ -74,14 +72,16 @@ def create_jl_obj(self) -> JuliaObj: def shape_type(self) -> tuple: return () + @dataclass(frozen=True) class DenseFormat(NestedLevelFormat): """Dense format wrapper type for Finch tensors. - + A dense format stores every slice of a tensor. A subfiber of a dense level is an array which stores every slice A[:, ..., :, i] as a distinct subfiber. Dense levels support both row-major and column-major access. """ + lvl: NestedLevelFormat dim_type: type = np.intp @@ -90,15 +90,17 @@ def create_jl_obj(self) -> JuliaObj: def shape_type(self) -> tuple: return self.lvl.shape_type + (self.dim_type,) - + + @dataclass(frozen=True) class SparseListFormat(NestedLevelFormat): """Sparse list format wrapper type for Finch tensors. - + A sparse list format stores only potentially non-fill slices using a sorted list. Slices that are entirely fill_value are omitted. This format is efficient for tensors with sparse patterns and supports column-major reads and bulk updates. """ + lvl: NestedLevelFormat dim_type: type = np.intp @@ -108,18 +110,20 @@ def create_jl_obj(self) -> JuliaObj: def shape_type(self) -> tuple: return self.lvl.shape_type + (self.dim_type,) + @dataclass(frozen=True) class SparseCOOFormat(NestedLevelFormat): """Coordinate (COO) format wrapper type for Finch tensors. - + A coordinate format stores sparse tensors as lists of coordinates. It uses N separate arrays to record which coordinates are stored, with coordinates sorted in column-major order. This is a legacy format maintained for backward compatibility. """ + lvl: NestedLevelFormat N: int = 2 - dim_type: Tuple{type} | None = np.intp + dim_type: tuple | None = np.intp def create_jl_obj(self) -> JuliaObj: return jl.SparseCOO(self.lvl.create_jl_obj()) @@ -129,53 +133,56 @@ def shape_type(self) -> tuple: return self.lvl.shape_type + (self.N * self.lvl.ndim,) return self.lvl.shape_type + self.dim_type + @dataclass(frozen=True) class SparseByteMapFormat(NestedLevelFormat): """Sparse byte map format wrapper type for Finch tensors. - + A sparse byte map format uses a dense bitmap to encode which slices are stored, similar to SparseList but supporting random access. Only potentially non-fill slices are stored as subfibers. """ + lvl: NestedLevelFormat dim_type: type = np.intp def create_jl_obj(self) -> JuliaObj: return jl.SparseByteMap(self.lvl.create_jl_obj()) - + def shape_type(self) -> tuple: return self.lvl.shape_type + (self.dim_type,) + # Helper Methods -def jlobj_to_format(obj: JuliaObj) -> LevelFType: +def jlobj_to_format(obj: JuliaObj) -> LevelFormat: """Construct a level hierarchy from a Julia Finch tensor object. - + Recursively constructs a Python representation of the tensor's level structure by inspecting the Julia object's levels. - + Parameters ---------- obj : JuliaObj A Julia Finch tensor object whose levels will be inspected. fill_value : float or int The fill value used for the tensor's sparse representation. - + Returns ------- - LevelFType + LevelFormat A Python representation of the level hierarchy. - + Raises ------ Exception If an unsupported level type is encountered. """ if jl.isa(obj, jl.Finch.Element): - return ElementLevel(jl.fill_value(obj)) + return ElementFormat(jl.fill_value(obj)) if jl.isa(obj, jl.Finch.Dense): - return DenseLevel(type(obj.shape), jlobj_to_format(obj.lvl)) + return DenseFormat(type(obj.shape), jlobj_to_format(obj.lvl)) if jl.isa(obj, jl.Finch.SparseList): - return SparseListLevel(type(obj.shape), jlobj_to_format(obj.lvl)) + return SparseListFormat(type(obj.shape), jlobj_to_format(obj.lvl)) if jl.isa(obj, jl.Finch.SparseByteMap): - return SparseByteMapLevel(jlobj_to_format(obj.lvl)) - raise Exception("Unhandled exception!") \ No newline at end of file + return SparseByteMapFormat(jlobj_to_format(obj.lvl)) + raise Exception("Unhandled exception!") diff --git a/src/finch/scalar.py b/src/finch/scalar.py deleted file mode 100644 index 47d4e8b..0000000 --- a/src/finch/scalar.py +++ /dev/null @@ -1,25 +0,0 @@ - -class ScalarFType(LevelFType): - def __init__(self, val: number): - self._val = val - - @property - def ndim(self) -> np.intp: - return np.intp(0) - - @property - def fill_value(self) -> Any: - return self._val - - @property - def element_type(self) -> Any: - return type(self._val) - - def __eq__(self, other): - return isinstance(other, ScalarFType) and self._val == other._val - - def __hash__(self): - return hash((self.__class__.__name__, self._val)) - - def create_jl_obj(self) -> JuliaObj: - return jl.Scalar(self._val) \ No newline at end of file diff --git a/src/finch/scheduler.py b/src/finch/scheduler.py index 2130b43..d50ad99 100644 --- a/src/finch/scheduler.py +++ b/src/finch/scheduler.py @@ -11,7 +11,7 @@ from finchlite.finch_logic import LogicLoader from .compiler import FinchJLCompiler -from .levels import DenseLevel, ElementLevel +from .levels import DenseFormat, ElementFormat from .tensor import FinchJLTensorFType @@ -23,9 +23,9 @@ def __init__( super().__init__(loader) def get_output_tns_ftype(self, fill_value: Any, shape_type: tuple[Any, ...]): - lvl = ElementLevel(fill_value) + lvl = ElementFormat(fill_value) for _ in shape_type: - lvl = DenseLevel(lvl) + lvl = DenseFormat(lvl) return FinchJLTensorFType(lvl) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 23d5754..d68692b 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -5,10 +5,14 @@ from finchlite import EagerTensor, Tensor, TensorFType from .julia import jc, jl -from .levels import LevelFType, ElementLevel, DenseLevel, SparseListLevel, SparseCOOLevel, jlobj_to_format -from .typing import JuliaObj, DType +from .levels import ( + LevelFormat, + jlobj_to_format, +) +from .typing import DType, JuliaObj from .utils import add_missing_dims, add_plus_one, expand_ellipsis + # Tensor Class and associated ftype class FinchJLTensorFType(TensorFType): def __init__(self, lvl): @@ -31,9 +35,6 @@ def shape_type(self) -> tuple[type, ...]: return self._lvl.shape_type def __call__(self, shape: tuple | None = None) -> Tensor: - if isinstance(self._lvl, Scalar): - return FinchJLTensor(self._lvl.create_jl_obj()) - if shape is None: raise ValueError("shape argument cannot be None for non scalar tensors.") return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), shape)) @@ -49,6 +50,7 @@ def __eq__(self, other): def __hash__(self): return hash(("FinchJLTensorFType", self._lvl)) + class FinchJLTensor(EagerTensor): def __init__(self, obj: JuliaObj): if isinstance(obj, JuliaObj): @@ -59,8 +61,6 @@ def __init__(self, obj: JuliaObj): @property def ftype(self) -> TensorFType: """Returns the ftype of the buffer""" - if self._is_scalar(): - return FinchJLTensorFType(Scalar(self._obj.val)) return FinchJLTensorFType(jlobj_to_format(self._obj, jl.fill_value(self._obj))) @property @@ -69,9 +69,6 @@ def shape(self) -> tuple: return jl.size(self._obj) def __getitem__(self, key): - if self._is_scalar(): - raise ValueError("Scalars are not subscriptable!") - if not isinstance(key, tuple): key = (key,) @@ -85,9 +82,6 @@ def __getitem__(self, key): return FinchJLTensor(result) return np.array(result) - def _is_scalar(self) -> bool: - return jl.isa(self._obj, jl.Finch.Scalar) - def _is_dense(self) -> bool: if self._is_scalar(): return False @@ -141,6 +135,7 @@ def __array_namespace__(self, *, api_version: str | None = None) -> Any: return finch + def asarray( obj, /, @@ -154,36 +149,35 @@ def asarray( if isinstance(obj, FinchJLTensor): if copy: return obj.copy() - else: - return obj - elif isinstance(obj, np.ndarray): + return obj + if isinstance(obj, np.ndarray): if copy: - if np.isfortran(obj): - obj = obj.copy() - else: - obj = np.asfortranarray(obj) + obj = obj.copy() if np.isfortran(obj) else np.asfortranarray(obj) else: if not np.isfortran(obj): - obj = np.asfortranarray(obj) - - lvl = ElementLevel(fill_value, NumpyBuffer(obj.reshape(-1))) + raise ValueError( + "Unable to avoid copy while creating an array as requested." + ) + + lvl = jl.ElementLevel(fill_value, obj.reshape(-1)) for i in obj.shape: lvl = jl.DenseLevel(lvl, i) return FinchJLTensor(lvl) - elif hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): + if hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): if obj.format == "coo": obj = obj.T if copy: if obj.format in ("coo", "csc"): - if not obj.has_sorted_indices: - obj = obj.sorted_indices() - else: - obj = obj.copy() + obj = obj.copy() if obj.has_sorted_indices else obj.sorted_indices() if not obj.has_canonical_format: obj.sum_duplicates() else: obj = obj.asformat("csc") - if copy is False and not obj.format in ("coo", "csc") and not obj.has_canonical_format: + if ( + copy is False + and obj.format not in ("coo", "csc") + and not obj.has_canonical_format + ): raise ValueError( "Unable to avoid copy while creating an array as requested." ) @@ -191,41 +185,28 @@ def asarray( if obj.format == "coo": return Tensor( jl.SparseCOOLevel( - jl.ElementLevel( - dtype, - fill_value, - obj.data - ), + (m, n), + jl.ElementLevel(dtype, fill_value, obj.data), 2, - idxs = ( + idxs=( jl.Finch.PlusOneVector(obj.cols), jl.Finch.PlusOneVector(obj.rows), ), - (m, n) ) ) - elif obj.format == "csc": + if obj.format == "csc": return Tensor( jl.DenseLevel( jl.SparseListLevel( - jl.ElementLevel( - dtype, - fill_value, - obj.data - ), + jl.ElementLevel(dtype, fill_value, obj.data), n, jl.Finch.PlusOneVector(obj.indptr), - jl.Finch.PlusOneVector(obj.indices) + jl.Finch.PlusOneVector(obj.indices), ), - m + m, ) ) - else: - raise ValueError(f"Unsupported SciPy format: {type(obj)}") - else: - raise ValueError( - "Either numpy array or a Finch tensor should " - f"be provided. Found: {type(obj)}" - ) - - + raise ValueError(f"Unsupported SciPy format: {type(obj)}") + raise ValueError( + f"Either numpy array or a Finch tensor should be provided. Found: {type(obj)}" + ) diff --git a/tests/conftest.py b/tests/conftest.py index 51d0cb4..455e172 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ import numpy as np + @pytest.fixture def rng(): return np.random.default_rng(42) diff --git a/tests/test_asarray.py b/tests/test_asarray.py index 7b1f6f3..5bc5998 100644 --- a/tests/test_asarray.py +++ b/tests/test_asarray.py @@ -1,12 +1,15 @@ """Tests for the asarray function.""" import pytest + import numpy as np + from finch import asarray from finch.tensor import FinchJLTensor try: import scipy.sparse as sp + HAS_SCIPY = True except ImportError: HAS_SCIPY = False @@ -179,13 +182,13 @@ def test_asarray_int64(self): def test_asarray_complex64(self): """Test converting complex64 array.""" - arr = np.array([[1+2j, 3+4j], [5+6j, 7+8j]], dtype=np.complex64) + arr = np.array([[1 + 2j, 3 + 4j], [5 + 6j, 7 + 8j]], dtype=np.complex64) result = asarray(arr) assert isinstance(result, FinchJLTensor) def test_asarray_complex128(self): """Test converting complex128 array.""" - arr = np.array([[1+2j, 3+4j], [5+6j, 7+8j]], dtype=np.complex128) + arr = np.array([[1 + 2j, 3 + 4j], [5 + 6j, 7 + 8j]], dtype=np.complex128) result = asarray(arr) assert isinstance(result, FinchJLTensor) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 510fede..1fde3ad 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -29,10 +29,10 @@ from finch.compiler import FinchJLCompiler, FinchJLKernel from finch.julia import jl -from finch.levels import DenseLevel, ElementLevel -from finch.tensor import FinchJLTensor +from finch.levels import DenseFormat, ElementFormat +from finch.tensor import FinchJLTensor, FinchJLTensorFType -a_format = DenseLevel(DenseLevel(ElementLevel(0))) +a_format = FinchJLTensorFType(DenseFormat(DenseFormat(ElementFormat(0)))) @pytest.mark.parametrize( diff --git a/tests/test_levels.py b/tests/test_levels.py deleted file mode 100644 index c4df3e5..0000000 --- a/tests/test_levels.py +++ /dev/null @@ -1,346 +0,0 @@ -"""Tests for the Finch levels module.""" - -import pytest -import numpy as np -from finch.levels import ( - ElementLevel, - DenseLevel, - PatternLevel, - SparseListLevel, - SparseByteMapLevel, - RepeatRLELevel, - SparseVBLLevel, - SparseCOOLevel, - SparseHashLevel, -) -from finch.buffer import NumpyBuffer - - -class TestElement: - """Test Element level construction.""" - - def test_element_creation_basic(self): - """Test creating an Element level with a fill value.""" - elem = ElementLevel(0.0) - assert elem._obj is not None - - def test_element_creation_with_int_fill(self): - """Test creating an Element level with integer fill value.""" - elem = ElementLevel(0) - assert elem._obj is not None - - def test_element_creation_with_float_fill(self): - """Test creating an Element level with float fill value.""" - elem = ElementLevel(3.14) - assert elem._obj is not None - - -class TestDense: - """Test Dense level construction.""" - - def test_dense_creation_basic(self): - """Test creating a Dense level.""" - elem = ElementLevel(0.0) - dense = DenseLevel(elem) - assert dense._obj is not None - - def test_dense_creation_with_shape(self): - """Test creating a Dense level with explicit shape.""" - elem = ElementLevel(0.0) - dense = DenseLevel(elem, shape=10) - assert dense._obj is not None - - def test_dense_nesting(self): - """Test creating nested Dense levels.""" - elem = ElementLevel(0.0) - dense1 = DenseLevel(elem) - dense2 = DenseLevel(dense1) - assert dense2._obj is not None - - -class TestPattern: - """Test Pattern level construction.""" - - def test_pattern_creation(self): - """Test creating a Pattern level.""" - pattern = PatternLevel() - assert pattern._obj is not None - - -class TestSparseList: - """Test SparseList level construction.""" - - def test_sparselist_creation_basic(self): - """Test creating a SparseList level.""" - elem = ElementLevel(0.0) - sparse = SparseListLevel(elem) - assert sparse._obj is not None - - def test_sparselist_creation_with_dim(self): - """Test creating a SparseList level with explicit dimension.""" - elem = ElementLevel(0.0) - sparse = SparseListLevel(elem, dim=10) - assert sparse._obj is not None - - def test_sparselist_creation_with_data_arrays(self): - """Test creating a SparseList level with pointer and index arrays.""" - elem = ElementLevel(0.0) - ptr = NumpyBuffer(np.array([0, 2, 2, 4], dtype=np.int32)) - idx = NumpyBuffer(np.array([1, 2, 1, 3], dtype=np.int32)) - sparse = SparseListLevel(elem, ptr=ptr, idx=idx) - assert sparse._obj is not None - - def test_sparselist_creation_with_data_lists(self): - """Test creating a SparseList level with pointer and index as lists.""" - elem = ElementLevel(0.0) - ptr = [0, 2, 2, 4] - idx = [1, 2, 1, 3] - sparse = SparseListLevel(elem, ptr=ptr, idx=idx) - assert sparse._obj is not None - - def test_sparselist_ptr_property(self): - """Test accessing ptr property of SparseList.""" - elem = ElementLevel(0.0) - sparse = SparseListLevel(elem) - # Property should be accessible - ptr_buffer = sparse.ptr - assert ptr_buffer is not None - - def test_sparselist_idx_property(self): - """Test accessing idx property of SparseList.""" - elem = ElementLevel(0.0) - sparse = SparseListLevel(elem) - # Property should be accessible - idx_buffer = sparse.idx - assert idx_buffer is not None - - -class TestSparseByteMap: - """Test SparseByteMap level construction.""" - - def test_sparsebytemap_creation_basic(self): - """Test creating a SparseByteMap level.""" - elem = ElementLevel(0.0) - sparse = SparseByteMapLevel(elem) - assert sparse._obj is not None - - def test_sparsebytemap_creation_with_dim(self): - """Test creating a SparseByteMap level with explicit dimension.""" - elem = ElementLevel(0.0) - sparse = SparseByteMapLevel(elem, dim=10) - assert sparse._obj is not None - - -class TestRepeatRLE: - """Test RepeatRLE level construction.""" - - def test_repeatrle_creation_basic(self): - """Test creating a RepeatRLE level.""" - elem = ElementLevel(0.0) - rle = RepeatRLELevel(elem) - assert rle._obj is not None - - def test_repeatrle_creation_with_dim(self): - """Test creating a RepeatRLE level with explicit dimension.""" - elem = ElementLevel(0.0) - rle = RepeatRLELevel(elem, dim=10) - assert rle._obj is not None - - -class TestSparseVBL: - """Test SparseVBL level construction.""" - - def test_sparsevbl_creation_basic(self): - """Test creating a SparseVBL level.""" - elem = ElementLevel(0.0) - vbl = SparseVBLLevel(elem) - assert vbl._obj is not None - - def test_sparsevbl_creation_with_dim(self): - """Test creating a SparseVBL level with explicit dimension.""" - elem = ElementLevel(0.0) - vbl = SparseVBLLevel(elem, dim=10) - assert vbl._obj is not None - - -class TestSparseCOO: - """Test SparseCOO level construction.""" - - def test_sparsecoo_creation_basic(self): - """Test creating a SparseCOO level.""" - elem = ElementLevel(0.0) - coo = SparseCOOLevel(2, elem) - assert coo._obj is not None - - def test_sparsecoo_creation_with_dims(self): - """Test creating a SparseCOO level with explicit dimensions.""" - elem = ElementLevel(0.0) - coo = SparseCOOLevel(2, elem, dims=(4, 3)) - assert coo._obj is not None - - def test_sparsecoo_creation_with_dims_list(self): - """Test creating a SparseCOO level with dimensions as list.""" - elem = ElementLevel(0.0) - coo = SparseCOOLevel(2, elem, dims=[4, 3]) - assert coo._obj is not None - - def test_sparsecoo_creation_with_coordinate_arrays(self): - """Test creating a SparseCOO level with coordinate arrays.""" - elem = ElementLevel(0.0) - i_coords = NumpyBuffer(np.array([0, 1, 2, 3], dtype=np.int32)) - j_coords = NumpyBuffer(np.array([0, 0, 2, 2], dtype=np.int32)) - coo = SparseCOOLevel(2, elem, tbl=(i_coords, j_coords)) - assert coo._obj is not None - - def test_sparsecoo_creation_with_coordinate_lists(self): - """Test creating a SparseCOO level with coordinate arrays as lists.""" - elem = ElementLevel(0.0) - i_coords = [0, 1, 2, 3] - j_coords = [0, 0, 2, 2] - coo = SparseCOOLevel(2, elem, tbl=(i_coords, j_coords)) - assert coo._obj is not None - - def test_sparsecoo_3d(self): - """Test creating a 3D SparseCOO level.""" - elem = ElementLevel(0.0) - coo = SparseCOOLevel(3, elem, dims=(5, 4, 3)) - assert coo._obj is not None - - def test_sparsecoo_tbl_property(self): - """Test accessing tbl property of SparseCOO.""" - elem = ElementLevel(0.0) - coo = SparseCOOLevel(2, elem) - # Property should be accessible - tbl = coo.tbl - assert tbl is not None - assert isinstance(tbl, tuple) - - -class TestSparseHash: - """Test SparseHash level construction.""" - - def test_sparsehash_creation_basic(self): - """Test creating a SparseHash level.""" - elem = ElementLevel(0.0) - hash_level = SparseHashLevel(2, elem) - assert hash_level._obj is not None - - def test_sparsehash_creation_with_dims(self): - """Test creating a SparseHash level with explicit dimensions.""" - elem = ElementLevel(0.0) - hash_level = SparseHashLevel(2, elem, dims=(4, 3)) - assert hash_level._obj is not None - - def test_sparsehash_creation_with_dims_list(self): - """Test creating a SparseHash level with dimensions as list.""" - elem = ElementLevel(0.0) - hash_level = SparseHashLevel(2, elem, dims=[4, 3]) - assert hash_level._obj is not None - - def test_sparsehash_3d(self): - """Test creating a 3D SparseHash level.""" - elem = ElementLevel(0.0) - hash_level = SparseHashLevel(3, elem, dims=(5, 4, 3)) - assert hash_level._obj is not None - - -class TestComposedLevels: - """Test composed level hierarchies.""" - - def test_csc_matrix_format(self): - """Test creating CSC matrix format (Dense(SparseList(Element))).""" - elem = ElementLevel(0.0) - sparse = SparseListLevel(elem) - dense = DenseLevel(sparse) - assert dense._obj is not None - - def test_csr_like_format(self): - """Test creating CSR-like format (SparseList(Dense(Element))).""" - elem = ElementLevel(0.0) - dense = DenseLevel(elem) - sparse = SparseListLevel(dense) - assert sparse._obj is not None - - def test_dcsc_format(self): - """Test creating DCSC format (SparseList(SparseList(Element))).""" - elem = ElementLevel(0.0) - sparse1 = SparseListLevel(elem) - sparse2 = SparseListLevel(sparse1) - assert sparse2._obj is not None - - def test_deep_nesting(self): - """Test deeply nested levels.""" - elem = ElementLevel(0.0) - sparse = SparseListLevel(elem) - dense = DenseLevel(sparse) - sparse2 = SparseListLevel(dense) - dense2 = DenseLevel(sparse2) - assert dense2._obj is not None - - -class TestEdgeCases: - """Test edge cases and error conditions.""" - - def test_element_with_negative_fill(self): - """Test Element with negative fill value.""" - elem = ElementLevel(-1.0) - assert elem._obj is not None - - def test_sparselist_only_ptr_no_idx(self): - """Test SparseList with ptr but no idx (should not add both).""" - elem = ElementLevel(0.0) - ptr = NumpyBuffer(np.array([0, 2, 2], dtype=np.int32)) - sparse = SparseListLevel(elem, ptr=ptr) - assert sparse._obj is not None - - def test_sparselist_only_idx_no_ptr(self): - """Test SparseList with idx but no ptr (should not add both).""" - elem = ElementLevel(0.0) - idx = NumpyBuffer(np.array([1, 2], dtype=np.int32)) - sparse = SparseListLevel(elem, idx=idx) - assert sparse._obj is not None - - def test_sparsecoo_single_coordinate(self): - """Test SparseCOO with single coordinate.""" - elem = ElementLevel(0.0) - coords = NumpyBuffer(np.array([0], dtype=np.int32)) - coo = SparseCOOLevel(1, elem, tbl=(coords,)) - assert coo._obj is not None - - def test_large_dimension(self): - """Test levels with large dimensions.""" - elem = ElementLevel(0.0) - sparse = SparseListLevel(elem, dim=1000000) - assert sparse._obj is not None - - -class TestArrayConversion: - """Test that array arguments are properly converted.""" - - def test_sparselist_converts_lists_to_arrays(self): - """Test that SparseList converts list arguments to arrays.""" - elem = ElementLevel(0.0) - ptr = NumpyBuffer(np.array([0, 2, 4], dtype=np.int32)) - idx = NumpyBuffer(np.array([1, 2, 3], dtype=np.int32)) - sparse = SparseListLevel(elem, ptr=ptr, idx=idx) - # Should not raise an error during creation - assert sparse._obj is not None - - def test_sparsecoo_converts_lists_to_arrays(self): - """Test that SparseCOO converts list arguments to arrays.""" - elem = ElementLevel(0.0) - coords_list = ( - NumpyBuffer(np.array([0, 1, 2], dtype=np.int32)), - NumpyBuffer(np.array([0, 1, 2], dtype=np.int32)) - ) - coo = SparseCOOLevel(2, elem, tbl=coords_list) - # Should not raise an error during creation - assert coo._obj is not None - - def test_different_dtype_arrays(self): - """Test that different dtype arrays are handled.""" - elem = ElementLevel(0.0) - ptr = NumpyBuffer(np.array([0, 2, 4], dtype=np.int64)) - idx = NumpyBuffer(np.array([1, 2, 3], dtype=np.int32)) - sparse = SparseListLevel(elem, ptr=ptr, idx=idx) - assert sparse._obj is not None From b3e24778e3dff6b750ce38d8ff8bf3b3890f26fa Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 17:32:14 -0400 Subject: [PATCH 52/81] row major --- src/finch/tensor.py | 50 +++++++++++++++++++-------------------------- 1 file changed, 21 insertions(+), 29 deletions(-) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index d68692b..ef41fe4 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -31,13 +31,11 @@ def element_type(self) -> Any: return self._lvl.element_type @property - def shape_type(self) -> tuple[type, ...]: - return self._lvl.shape_type + def shape_type(self) -> tuple: + return reversed(self._lvl.shape_type) - def __call__(self, shape: tuple | None = None) -> Tensor: - if shape is None: - raise ValueError("shape argument cannot be None for non scalar tensors.") - return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), shape)) + def __call__(self, shape: tuple) -> Tensor: + return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), reversed(shape))) def from_numpy(self, _) -> Tensor: raise NotImplementedError @@ -77,15 +75,10 @@ def __getitem__(self, key): key = add_missing_dims(key, self.shape) key = add_plus_one(key, self.shape) - result = self._obj[key] - if jl.isa(result, jl.Finch.Tensor): - return FinchJLTensor(result) - return np.array(result) + result = self._obj[reversed(key)] + return FinchJLTensor(jl.Tensor(result)) def _is_dense(self) -> bool: - if self._is_scalar(): - return False - lvl = self._obj.lvl for _ in self.shape: if not jl.isa(lvl, jl.Finch.Dense): @@ -94,9 +87,6 @@ def _is_dense(self) -> bool: return True def todense(self) -> np.ndarray: - if self._is_scalar(): - return np.asarray(self._obj.val) - obj = self._obj if self._is_dense: @@ -114,16 +104,19 @@ def todense(self) -> np.ndarray: for _ in range(self.ndim): dense_tensor = dense_tensor.lvl - return np.asarray(jl.reshape(dense_tensor.val, shape)) + arr = jl.reshape(dense_tensor.val) + return np.asarray(np.permute_dims(arr, reversed(range(self.ndims)))) def __eq__(self, other): return isinstance(other, FinchJLTensor) and self._obj == other._obj def __repr__(self): - return jl.sprint(jl.show, self._obj) + swiz = jl.swizzle(self._obj, reversed(range(self.ndim,1,-1))) + return jl.sprint(jl.show, swiz) def __str__(self): - return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) + swiz = jl.swizzle(self._obj, reversed(range(self.ndim,1,-1))) + return jl.sprint(jl.show, jl.MIME("text/plain"), swiz) def __array_namespace__(self, *, api_version: str | None = None) -> Any: if api_version is None: @@ -158,24 +151,23 @@ def asarray( raise ValueError( "Unable to avoid copy while creating an array as requested." ) + buf = np.reshape(np.permute_dims(obj, reversed(range(obj.ndim))), -1) - lvl = jl.ElementLevel(fill_value, obj.reshape(-1)) + lvl = jl.ElementLevel(fill_value, buf) for i in obj.shape: lvl = jl.DenseLevel(lvl, i) return FinchJLTensor(lvl) if hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): - if obj.format == "coo": - obj = obj.T if copy: - if obj.format in ("coo", "csc"): + if obj.format in ("coo", "csr"): obj = obj.copy() if obj.has_sorted_indices else obj.sorted_indices() if not obj.has_canonical_format: obj.sum_duplicates() else: - obj = obj.asformat("csc") + obj = obj.asformat("csr") if ( copy is False - and obj.format not in ("coo", "csc") + and obj.format not in ("coo", "csr") and not obj.has_canonical_format ): raise ValueError( @@ -185,7 +177,7 @@ def asarray( if obj.format == "coo": return Tensor( jl.SparseCOOLevel( - (m, n), + (n, m), jl.ElementLevel(dtype, fill_value, obj.data), 2, idxs=( @@ -194,16 +186,16 @@ def asarray( ), ) ) - if obj.format == "csc": + if obj.format == "csr": return Tensor( jl.DenseLevel( jl.SparseListLevel( jl.ElementLevel(dtype, fill_value, obj.data), - n, + m, jl.Finch.PlusOneVector(obj.indptr), jl.Finch.PlusOneVector(obj.indices), ), - m, + n, ) ) raise ValueError(f"Unsupported SciPy format: {type(obj)}") From 9cf5967ea530c6289a8931e32f50e706b5ca513c Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 18:00:43 -0400 Subject: [PATCH 53/81] cool --- src/finch/levels.py | 10 ++++++---- src/finch/tensor.py | 31 ++++++++++++++++++++++++------- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 38063cc..53517b0 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -178,11 +178,13 @@ def jlobj_to_format(obj: JuliaObj) -> LevelFormat: If an unsupported level type is encountered. """ if jl.isa(obj, jl.Finch.Element): - return ElementFormat(jl.fill_value(obj)) + return ElementFormat(jl.Finch.level_fill_value(jl.typeof(obj))) if jl.isa(obj, jl.Finch.Dense): - return DenseFormat(type(obj.shape), jlobj_to_format(obj.lvl)) + return DenseFormat(jlobj_to_format(obj.lvl), type(obj.shape)) if jl.isa(obj, jl.Finch.SparseList): - return SparseListFormat(type(obj.shape), jlobj_to_format(obj.lvl)) + return SparseListFormat(jlobj_to_format(obj.lvl), type(obj.shape)) + if jl.isa(obj, jl.Finch.SparseCOO): + return SparseCOOFormat(jlobj_to_format(obj.lvl), type(obj).__type_params__[0], type(obj.shape)) if jl.isa(obj, jl.Finch.SparseByteMap): - return SparseByteMapFormat(jlobj_to_format(obj.lvl)) + return SparseByteMapFormat(jlobj_to_format(obj.lvl), type(obj.shape)) raise Exception("Unhandled exception!") diff --git a/src/finch/tensor.py b/src/finch/tensor.py index ef41fe4..ac7b6b2 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -30,6 +30,10 @@ def fill_value(self) -> Any: def element_type(self) -> Any: return self._lvl.element_type + @property + def dtype(self) -> Any: + return self.element_type + @property def shape_type(self) -> tuple: return reversed(self._lvl.shape_type) @@ -52,6 +56,7 @@ def __hash__(self): class FinchJLTensor(EagerTensor): def __init__(self, obj: JuliaObj): if isinstance(obj, JuliaObj): + assert jl.isa(obj, jl.Finch.Tensor) self._obj = obj else: raise ValueError(f"Raw julia object expected. Found: {type(obj)}") @@ -59,7 +64,11 @@ def __init__(self, obj: JuliaObj): @property def ftype(self) -> TensorFType: """Returns the ftype of the buffer""" - return FinchJLTensorFType(jlobj_to_format(self._obj, jl.fill_value(self._obj))) + return FinchJLTensorFType(jlobj_to_format(self._obj.lvl)) + + @property + def dtype(self) -> Any: + return self.element_type @property def shape(self) -> tuple: @@ -139,10 +148,18 @@ def asarray( ) -> FinchJLTensor: if fill_value is None: fill_value = 0.0 + if copy is None: + copy = True if isinstance(obj, FinchJLTensor): if copy: return obj.copy() return obj + if isinstance(obj, int | float | complex | bool | list): + if copy is False: + raise ValueError( + "copy=False isn't supported for scalar inputs and Python lists" + ) + obj = np.asarray(obj) if isinstance(obj, np.ndarray): if copy: obj = obj.copy() if np.isfortran(obj) else np.asfortranarray(obj) @@ -151,12 +168,12 @@ def asarray( raise ValueError( "Unable to avoid copy while creating an array as requested." ) - buf = np.reshape(np.permute_dims(obj, reversed(range(obj.ndim))), -1) + buf = np.reshape(np.permute_dims(obj, tuple(reversed(range(obj.ndim)))), -1) lvl = jl.ElementLevel(fill_value, buf) for i in obj.shape: lvl = jl.DenseLevel(lvl, i) - return FinchJLTensor(lvl) + return FinchJLTensor(jl.Tensor(lvl)) if hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): if copy: if obj.format in ("coo", "csr"): @@ -175,7 +192,7 @@ def asarray( ) m, n = obj.shape if obj.format == "coo": - return Tensor( + return FinchJLTensor(jl.Tensor( jl.SparseCOOLevel( (n, m), jl.ElementLevel(dtype, fill_value, obj.data), @@ -185,9 +202,9 @@ def asarray( jl.Finch.PlusOneVector(obj.rows), ), ) - ) + )) if obj.format == "csr": - return Tensor( + return FinchJLTensor(jl.Tensor( jl.DenseLevel( jl.SparseListLevel( jl.ElementLevel(dtype, fill_value, obj.data), @@ -197,7 +214,7 @@ def asarray( ), n, ) - ) + )) raise ValueError(f"Unsupported SciPy format: {type(obj)}") raise ValueError( f"Either numpy array or a Finch tensor should be provided. Found: {type(obj)}" From ad5df8d3d9d639c16202e7cc1caf145b73116d50 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 18:31:41 -0400 Subject: [PATCH 54/81] fix --- src/finch/levels.py | 19 ++++++++++++++----- src/finch/tensor.py | 20 ++++++++++++-------- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index 53517b0..fa935c8 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -45,8 +45,9 @@ class ElementFormat(LevelFormat): The element level is a leaf level used at the end of the tensor tree structure. """ - def __init__(self, fill_value: number): + def __init__(self, fill_value: number, element_type: Any | None = None): self._fill_value = fill_value + self._element_type = type(fill_value) if element_type is None else element_type @property def ndim(self) -> np.intp: @@ -58,13 +59,17 @@ def fill_value(self) -> Any: @property def element_type(self) -> Any: - return type(self._fill_value) + return self._element_type def __eq__(self, other): - return isinstance(other, ElementFormat) and self._fill_value == other.fill_value + return ( + isinstance(other, ElementFormat) + and self._fill_value == other.fill_value + and self._element_type == other.element_type + ) def __hash__(self): - return hash((self.__class__.__name__, self._fill_value)) + return hash((self.__class__.__name__, self._fill_value, self._element_type)) def create_jl_obj(self) -> JuliaObj: return jl.Element(self._fill_value) @@ -178,7 +183,11 @@ def jlobj_to_format(obj: JuliaObj) -> LevelFormat: If an unsupported level type is encountered. """ if jl.isa(obj, jl.Finch.Element): - return ElementFormat(jl.Finch.level_fill_value(jl.typeof(obj))) + obj_type = jl.typeof(obj) + return ElementFormat( + jl.Finch.level_fill_value(obj_type), + jl.Finch.level_eltype(obj_type), + ) if jl.isa(obj, jl.Finch.Dense): return DenseFormat(jlobj_to_format(obj.lvl), type(obj.shape)) if jl.isa(obj, jl.Finch.SparseList): diff --git a/src/finch/tensor.py b/src/finch/tensor.py index ac7b6b2..4bb1cf3 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -36,10 +36,12 @@ def dtype(self) -> Any: @property def shape_type(self) -> tuple: - return reversed(self._lvl.shape_type) + return tuple(reversed(self._lvl.shape_type)) def __call__(self, shape: tuple) -> Tensor: - return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), reversed(shape))) + return FinchJLTensor( + jl.Finch.Tensor(self._lvl.create_jl_obj(), tuple(reversed(shape))) + ) def from_numpy(self, _) -> Tensor: raise NotImplementedError @@ -84,8 +86,10 @@ def __getitem__(self, key): key = add_missing_dims(key, self.shape) key = add_plus_one(key, self.shape) - result = self._obj[reversed(key)] - return FinchJLTensor(jl.Tensor(result)) + result = self._obj[tuple(reversed(key))] + if jl.isa(result, jl.Finch.Tensor): + return FinchJLTensor(result) + return result def _is_dense(self) -> bool: lvl = self._obj.lvl @@ -113,18 +117,18 @@ def todense(self) -> np.ndarray: for _ in range(self.ndim): dense_tensor = dense_tensor.lvl - arr = jl.reshape(dense_tensor.val) - return np.asarray(np.permute_dims(arr, reversed(range(self.ndims)))) + arr = jl.reshape(dense_tensor.val, tuple(shape)) + return np.asarray(arr) def __eq__(self, other): return isinstance(other, FinchJLTensor) and self._obj == other._obj def __repr__(self): - swiz = jl.swizzle(self._obj, reversed(range(self.ndim,1,-1))) + swiz = jl.swizzle(self._obj, tuple(reversed(range(self.ndim, 1, -1)))) return jl.sprint(jl.show, swiz) def __str__(self): - swiz = jl.swizzle(self._obj, reversed(range(self.ndim,1,-1))) + swiz = jl.swizzle(self._obj, tuple(reversed(range(self.ndim, 1, -1)))) return jl.sprint(jl.show, jl.MIME("text/plain"), swiz) def __array_namespace__(self, *, api_version: str | None = None) -> Any: From a4467c412dc6f421b2570bd13d1181c0a69f08a9 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 18:33:58 -0400 Subject: [PATCH 55/81] fix --- src/finch/levels.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index fa935c8..c2c8a4b 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -22,7 +22,7 @@ def ndim(self) -> np.intp: @property def fill_value(self) -> Any: - return self.lvl.fill_value + return self.element_type(self.lvl.fill_value) @property def element_type(self) -> Any: @@ -55,7 +55,7 @@ def ndim(self) -> np.intp: @property def fill_value(self) -> Any: - return self._fill_value + return self.element_type(self._fill_value) @property def element_type(self) -> Any: From 2a9b9299f20678c6d449bdbfa6da0b3072e25b52 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 18:47:01 -0400 Subject: [PATCH 56/81] fix --- README.md | 26 +++++++++++++++++++++++ tests/conftest.py | 13 ++++++++++++ tests/test_array_api.py | 47 ++++++++++++++++++++++++++++++++++++++--- 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 84b82cb..ae280c0 100644 --- a/README.md +++ b/README.md @@ -74,3 +74,29 @@ tests: ```bash poetry run pytest ``` + +Array API tests are included in `tests/test_array_api.py`. These tests invoke the [Array API Conformance Tests](https://github.com/data-apis/array-api-tests). +To forward `pytest` options to the nested +`array-api-tests` invocation, use `--array-api` (alias: +`--array-api-pytest-args`): + +```bash +poetry run pytest tests/test_array_api.py \ + --array-api="-k creation_functions" \ + --array-api="-x" +``` + +By default, the nested Array API run forwards common top-level pytest options +from your main invocation: + +- `-x`/`--maxfail` +- `-s` +- `-v`, `-vv`, etc. +- `-k` + +You can repeat `--array-api` (or `--array-api-pytest-args`) multiple times. +Each value is parsed like shell arguments and appended to the nested `pytest` +call. + +`ARRAY_API_TESTS_ARGS` is still supported as a fallback for compatibility, but +the CLI option is preferred. diff --git a/tests/conftest.py b/tests/conftest.py index 455e172..094fee3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,6 +3,19 @@ import numpy as np +def pytest_addoption(parser): + parser.addoption( + "--array-api", + "--array-api-pytest-args", + action="append", + default=[], + help=( + "Arguments forwarded to the nested array-api-tests pytest run. " + "Repeat this option to pass multiple groups." + ), + ) + + @pytest.fixture def rng(): return np.random.default_rng(42) diff --git a/tests/test_array_api.py b/tests/test_array_api.py index b6440f1..9d66a88 100644 --- a/tests/test_array_api.py +++ b/tests/test_array_api.py @@ -1,9 +1,46 @@ import os +import shlex import subprocess import sys -def test_array_api(): +def _get_forwarded_main_pytest_args(request): + args = [] + + maxfail = request.config.getoption("maxfail") + if maxfail: + args.append(f"--maxfail={maxfail}") + + capture = request.config.getoption("capture") + if capture == "no": + args.append("-s") + + verbose = request.config.getoption("verbose") + if verbose and verbose > 0: + args.append(f"-{'v' * verbose}") + + keyword = request.config.getoption("keyword") + if keyword: + args.extend(["-k", keyword]) + + return args + + +def _get_user_array_api_pytest_args(request): + cli_args = [] + for arg_group in request.config.getoption("array_api_pytest_args"): + cli_args.extend(shlex.split(arg_group)) + if cli_args: + return cli_args + + env_args = os.environ.get("ARRAY_API_TESTS_ARGS") + if env_args: + return shlex.split(env_args) + + return ["-vv", "-s"] + + +def test_array_api(request): ARRAY_API_TESTS_DIR = os.environ.get( "ARRAY_API_TESTS_DIR", os.path.abspath( @@ -19,7 +56,9 @@ def test_array_api(): os.path.join(os.path.dirname(__file__), "../array-api-skips.txt"), ), ) - ARRAY_API_TESTS_ARGS = os.environ.get("ARRAY_API_TESTS_ARGS", "-vv -s") + FORWARDED_MAIN_PYTEST_ARGS = _get_forwarded_main_pytest_args(request) + ARRAY_API_TESTS_ARGS = _get_user_array_api_pytest_args(request) + NESTED_PYTEST_ARGS = [*FORWARDED_MAIN_PYTEST_ARGS, *ARRAY_API_TESTS_ARGS] print(f"[array-api] using dir: {ARRAY_API_TESTS_DIR}", flush=True) print(f"[array-api] target rev: {ARRAY_API_TESTS_REV}", flush=True) @@ -81,12 +120,14 @@ def test_array_api(): # Run the tests using pytest print("[array-api] running external array-api-tests...", flush=True) + print(f"[array-api] forwarded main pytest args: {FORWARDED_MAIN_PYTEST_ARGS}", flush=True) + print(f"[array-api] user nested pytest args: {ARRAY_API_TESTS_ARGS}", flush=True) result = subprocess.run( [ sys.executable, "-m", "pytest", - *ARRAY_API_TESTS_ARGS.split(), + *NESTED_PYTEST_ARGS, f"{ARRAY_API_TESTS_DIR}/array_api_tests/", "--max-examples=2", "--derandomize", From ae8eb1106d52f60d3976b6a14fa52bc255ea928c Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 18:50:31 -0400 Subject: [PATCH 57/81] fix --- tests/test_array_api.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_array_api.py b/tests/test_array_api.py index 9d66a88..a75082c 100644 --- a/tests/test_array_api.py +++ b/tests/test_array_api.py @@ -120,7 +120,10 @@ def test_array_api(request): # Run the tests using pytest print("[array-api] running external array-api-tests...", flush=True) - print(f"[array-api] forwarded main pytest args: {FORWARDED_MAIN_PYTEST_ARGS}", flush=True) + print( + f"[array-api] forwarded main pytest args: {FORWARDED_MAIN_PYTEST_ARGS}", + flush=True, + ) print(f"[array-api] user nested pytest args: {ARRAY_API_TESTS_ARGS}", flush=True) result = subprocess.run( [ From 7dc056989d0606759a1c4423fd736ab9cfc34efb Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 18:53:54 -0400 Subject: [PATCH 58/81] fix --- tests/conftest.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/conftest.py b/tests/conftest.py index 094fee3..61800c9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,6 +7,7 @@ def pytest_addoption(parser): parser.addoption( "--array-api", "--array-api-pytest-args", + dest="array_api_pytest_args", action="append", default=[], help=( From 9f35a2729a0421a14a8ee09b65f9b74ae704dccc Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 18:54:41 -0400 Subject: [PATCH 59/81] fix --- src/finch/levels.py | 4 +++- src/finch/tensor.py | 44 ++++++++++++++++++++++++-------------------- 2 files changed, 27 insertions(+), 21 deletions(-) diff --git a/src/finch/levels.py b/src/finch/levels.py index c2c8a4b..03eb8bf 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -193,7 +193,9 @@ def jlobj_to_format(obj: JuliaObj) -> LevelFormat: if jl.isa(obj, jl.Finch.SparseList): return SparseListFormat(jlobj_to_format(obj.lvl), type(obj.shape)) if jl.isa(obj, jl.Finch.SparseCOO): - return SparseCOOFormat(jlobj_to_format(obj.lvl), type(obj).__type_params__[0], type(obj.shape)) + return SparseCOOFormat( + jlobj_to_format(obj.lvl), type(obj).__type_params__[0], type(obj.shape) + ) if jl.isa(obj, jl.Finch.SparseByteMap): return SparseByteMapFormat(jlobj_to_format(obj.lvl), type(obj.shape)) raise Exception("Unhandled exception!") diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 4bb1cf3..562c80d 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -196,29 +196,33 @@ def asarray( ) m, n = obj.shape if obj.format == "coo": - return FinchJLTensor(jl.Tensor( - jl.SparseCOOLevel( - (n, m), - jl.ElementLevel(dtype, fill_value, obj.data), - 2, - idxs=( - jl.Finch.PlusOneVector(obj.cols), - jl.Finch.PlusOneVector(obj.rows), - ), + return FinchJLTensor( + jl.Tensor( + jl.SparseCOOLevel( + (n, m), + jl.ElementLevel(dtype, fill_value, obj.data), + 2, + idxs=( + jl.Finch.PlusOneVector(obj.cols), + jl.Finch.PlusOneVector(obj.rows), + ), + ) ) - )) + ) if obj.format == "csr": - return FinchJLTensor(jl.Tensor( - jl.DenseLevel( - jl.SparseListLevel( - jl.ElementLevel(dtype, fill_value, obj.data), - m, - jl.Finch.PlusOneVector(obj.indptr), - jl.Finch.PlusOneVector(obj.indices), - ), - n, + return FinchJLTensor( + jl.Tensor( + jl.DenseLevel( + jl.SparseListLevel( + jl.ElementLevel(dtype, fill_value, obj.data), + m, + jl.Finch.PlusOneVector(obj.indptr), + jl.Finch.PlusOneVector(obj.indices), + ), + n, + ) ) - )) + ) raise ValueError(f"Unsupported SciPy format: {type(obj)}") raise ValueError( f"Either numpy array or a Finch tensor should be provided. Found: {type(obj)}" From ccdc6e2700bc8e231a2ec72a33e7734697930645 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 19:26:28 -0400 Subject: [PATCH 60/81] fix --- src/finch/__init__.py | 22 +++++++ src/finch/tensor.py | 142 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 164 insertions(+) diff --git a/src/finch/__init__.py b/src/finch/__init__.py index f7f916a..d2d2588 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -124,7 +124,18 @@ from .tensor import ( FinchJLTensor, FinchJLTensorFType, + arange, asarray, + empty, + empty_like, + full, + full_like, + linspace, + ones, + ones_like, + reshape, + zeros, + zeros_like, ) __all__ = [ @@ -138,6 +149,7 @@ "add", "all", "any", + "arange", "asarray", "asin", "asinh", @@ -168,6 +180,8 @@ "einop", "einsum", "elementwise", + "empty", + "empty_like", "equal", "exp", "expand_dims", @@ -179,6 +193,8 @@ "float64", "floor", "floordiv", + "full", + "full_like", "fuse", "fused", "get_default_scheduler", @@ -198,6 +214,7 @@ "lazy", "less", "less_equal", + "linspace", "log", "log1p", "log2", @@ -220,6 +237,8 @@ "negative", "nextafter", "not_equal", + "ones", + "ones_like", "permute_dims", "positive", "pow", @@ -229,6 +248,7 @@ "reciprocal", "reduce", "remainder", + "reshape", "round", "set_default_scheduler", "sign", @@ -255,4 +275,6 @@ "uint64", "var", "vecdot", + "zeros", + "zeros_like", ] diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 562c80d..f193272 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -4,9 +4,12 @@ from finchlite import EagerTensor, Tensor, TensorFType +from . import dtypes as jl_dtypes from .julia import jc, jl from .levels import ( + ElementFormat, LevelFormat, + SparseCOOFormat, jlobj_to_format, ) from .typing import DType, JuliaObj @@ -227,3 +230,142 @@ def asarray( raise ValueError( f"Either numpy array or a Finch tensor should be provided. Found: {type(obj)}" ) + + +def reshape( + x: FinchJLTensor, /, shape: tuple[int, ...], *, copy: bool | None = None +) -> FinchJLTensor: + if copy is False: + raise ValueError("Unable to avoid copy during reshape.") + if all(i == 1 for i in x.shape): + return full(shape, x[()], dtype=x.dtype) + return FinchJLTensor(jl.reshape(x._obj, tuple(reversed(shape)))) + + +def full( + shape: int | tuple[int, ...], + val: jl_dtypes.number, + *, + dtype: DType | None = None, + format=None, +) -> FinchJLTensor: + if not np.isscalar(val): + raise ValueError("`fill_value` must be a scalar") + if isinstance(shape, int): + shape = (shape,) + dtype = ( + np.asarray(val).dtype.type if dtype is None else jl_dtypes.jl_to_np_dtype[dtype] + ) + if dtype == np.bool_: # Fails with: Finch currently only supports isbits defaults + dtype = bool + + if format is None: + format = SparseCOOFormat(ElementFormat(val, dtype), len(shape)) + + if format.fill_value != val: + return FinchJLTensor( + jl.Tensor(format.construct_julia_lvl(), np.full(val, reversed(shape))) + ) + return FinchJLTensor(jl.Tensor(format.construct_julia_lvl(), *reversed(shape))) + + +def full_like( + x: FinchJLTensor, + /, + fill_value: jl_dtypes.number, + *, + dtype: DType | None = None, + format: str = "coo", +) -> FinchJLTensor: + return full(x.shape, fill_value, dtype=dtype, format=format) + + +def ones( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + format: str = "coo", +) -> FinchJLTensor: + return full(shape, np.float64(1), dtype=dtype, format=format) + + +def ones_like( + x: FinchJLTensor, + /, + *, + dtype: DType | None = None, + format: str = "coo", +) -> FinchJLTensor: + dtype = x.dtype if dtype is None else dtype + return ones(x.shape, dtype=dtype, format=format) + + +def zeros( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + format: str = "coo", +) -> FinchJLTensor: + return full(shape, np.float64(0), dtype=dtype, format=format) + + +def zeros_like( + x: FinchJLTensor, + /, + *, + dtype: DType | None = None, + format: str = "coo", +) -> FinchJLTensor: + dtype = x.dtype if dtype is None else dtype + return zeros(x.shape, dtype=dtype, format=format) + + +def empty( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + format: str = "coo", +) -> FinchJLTensor: + return full(shape, np.float64(0), dtype=dtype, format=format) + + +def empty_like( + x: FinchJLTensor, + /, + *, + dtype: DType | None = None, + format: str = "coo", +) -> FinchJLTensor: + dtype = x.dtype if dtype is None else dtype + return empty(x.shape, dtype=dtype, format=format) + + +def arange( + start: float, + /, + stop: float | None = None, + step: float = 1, + *, + dtype: DType | None = None, +) -> FinchJLTensor: + return asarray(np.arange(start, stop, step, jl_dtypes.jl_to_np_dtype[dtype])) + + +def linspace( + start: complex, + stop: complex, + /, + num: int, + *, + dtype: DType | None = None, + endpoint: bool = True, +) -> FinchJLTensor: + return asarray( + np.linspace( + start, + stop, + num=num, + dtype=jl_dtypes.jl_to_np_dtype[dtype], + endpoint=endpoint, + ) + ) From 0da6e875ba790effa3fe0dabe0708c62c139ad77 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 20:11:48 -0400 Subject: [PATCH 61/81] fixing --- src/finch/tensor.py | 18 +++++++++--------- src/finch/utils.py | 2 -- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index f193272..99b262d 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -89,10 +89,11 @@ def __getitem__(self, key): key = add_missing_dims(key, self.shape) key = add_plus_one(key, self.shape) - result = self._obj[tuple(reversed(key))] - if jl.isa(result, jl.Finch.Tensor): - return FinchJLTensor(result) - return result + result = jl.getindex(self._obj, *reversed(key)) + if all(isinstance(k, int) for k in key): + return result + assert jl.isa(result, jl.Finch.Tensor) + return FinchJLTensor(result) def _is_dense(self) -> bool: lvl = self._obj.lvl @@ -127,8 +128,7 @@ def __eq__(self, other): return isinstance(other, FinchJLTensor) and self._obj == other._obj def __repr__(self): - swiz = jl.swizzle(self._obj, tuple(reversed(range(self.ndim, 1, -1)))) - return jl.sprint(jl.show, swiz) + return jl.sprint(jl.show, self._obj) def __str__(self): swiz = jl.swizzle(self._obj, tuple(reversed(range(self.ndim, 1, -1)))) @@ -238,7 +238,7 @@ def reshape( if copy is False: raise ValueError("Unable to avoid copy during reshape.") if all(i == 1 for i in x.shape): - return full(shape, x[()], dtype=x.dtype) + return full(shape, x[tuple(i - 1 for i in x.shape)], dtype=x.dtype) return FinchJLTensor(jl.reshape(x._obj, tuple(reversed(shape)))) @@ -264,9 +264,9 @@ def full( if format.fill_value != val: return FinchJLTensor( - jl.Tensor(format.construct_julia_lvl(), np.full(val, reversed(shape))) + jl.Tensor(format.create_jl_obj(), np.full(val, reversed(shape))) ) - return FinchJLTensor(jl.Tensor(format.construct_julia_lvl(), *reversed(shape))) + return FinchJLTensor(jl.Tensor(format.create_jl_obj(), *reversed(shape))) def full_like( diff --git a/src/finch/utils.py b/src/finch/utils.py index acdf3fb..3afa309 100644 --- a/src/finch/utils.py +++ b/src/finch/utils.py @@ -53,8 +53,6 @@ def _slice_plus_one(s: slice, size: int) -> range: else: stop = stop_default - if (start, stop, step) == (1, size, 1): - return jl.Colon() return jl.range(start=start, step=step, stop=stop) From ca3d3534d09dcdd0ee81986fede7eb3cd9c32cc7 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 20:21:39 -0400 Subject: [PATCH 62/81] fix --- src/finch/__init__.py | 13 +++++++++++++ src/finch/tensor.py | 38 +++++++++++++++++++++++++++++++++++++- 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/src/finch/__init__.py b/src/finch/__init__.py index d2d2588..b79b7ea 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -1,3 +1,5 @@ +import math + from finchlite import ( abs, acos, @@ -138,6 +140,12 @@ zeros_like, ) +e = math.e +pi = math.pi +inf = math.inf +nan = math.nan +newaxis = None + __all__ = [ "COMPILE_JULIA", "FinchJLTensor", @@ -180,6 +188,7 @@ "einop", "einsum", "elementwise", + "e", "empty", "empty_like", "equal", @@ -203,6 +212,7 @@ "hypot", "iinfo", "imag", + "inf", "int8", "int16", "int32", @@ -234,12 +244,15 @@ "mod", "moveaxis", "multiply", + "nan", "negative", "nextafter", + "newaxis", "not_equal", "ones", "ones_like", "permute_dims", + "pi", "positive", "pow", "power", diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 99b262d..01dc0df 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,4 +1,5 @@ from typing import Any +import operator import numpy as np @@ -84,6 +85,10 @@ def __getitem__(self, key): if not isinstance(key, tuple): key = (key,) + # Array API behavior: indexing a 0-D array with () is a no-op. + if self.shape == () and key == (): + return self + # standard indexing mode key = expand_ellipsis(key, self.shape) key = add_missing_dims(key, self.shape) @@ -144,6 +149,28 @@ def __array_namespace__(self, *, api_version: str | None = None) -> Any: return finch + def _scalar_value(self): + if self.shape != (): + raise TypeError("only 0-dimensional arrays can be converted to Python scalars") + return jl.getindex(self._obj) + + def __bool__(self) -> bool: + return bool(self._scalar_value()) + + def __int__(self) -> int: + return int(self._scalar_value()) + + def __float__(self) -> float: + return float(self._scalar_value()) + + def __index__(self) -> int: + try: + return operator.index(self._scalar_value()) + except TypeError as exc: + raise TypeError( + "only integer scalar arrays can be converted to an index" + ) from exc + def asarray( obj, @@ -168,6 +195,8 @@ def asarray( ) obj = np.asarray(obj) if isinstance(obj, np.ndarray): + if dtype is not None: + obj = np.asarray(obj, dtype=jl_dtypes.jl_to_np_dtype[dtype]) if copy: obj = obj.copy() if np.isfortran(obj) else np.asfortranarray(obj) else: @@ -177,7 +206,7 @@ def asarray( ) buf = np.reshape(np.permute_dims(obj, tuple(reversed(range(obj.ndim)))), -1) - lvl = jl.ElementLevel(fill_value, buf) + lvl = jl.ElementLevel(np.asarray(fill_value, dtype=obj.dtype).item(), buf) for i in obj.shape: lvl = jl.DenseLevel(lvl, i) return FinchJLTensor(jl.Tensor(lvl)) @@ -198,6 +227,8 @@ def asarray( "Unable to avoid copy while creating an array as requested." ) m, n = obj.shape + if dtype is not None: + fill_value = np.asarray(fill_value, dtype=jl_dtypes.jl_to_np_dtype[dtype]).item() if obj.format == "coo": return FinchJLTensor( jl.Tensor( @@ -259,6 +290,11 @@ def full( if dtype == np.bool_: # Fails with: Finch currently only supports isbits defaults dtype = bool + # Rank-0 tensors should be represented as a leaf element level. + # Building them through SparseCOO requires an explicit rank parameter. + if len(shape) == 0 and format is None: + return FinchJLTensor(jl.Tensor(ElementFormat(val, dtype).create_jl_obj())) + if format is None: format = SparseCOOFormat(ElementFormat(val, dtype), len(shape)) From 34b7ad2014aecd9557480fb516a0afefd7b82a3e Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 27 Mar 2026 20:30:02 -0400 Subject: [PATCH 63/81] fix --- src/finch/levels.py | 5 +++++ src/finch/tensor.py | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/src/finch/levels.py b/src/finch/levels.py index 03eb8bf..9f9f72a 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -74,6 +74,7 @@ def __hash__(self): def create_jl_obj(self) -> JuliaObj: return jl.Element(self._fill_value) + @property def shape_type(self) -> tuple: return () @@ -93,6 +94,7 @@ class DenseFormat(NestedLevelFormat): def create_jl_obj(self) -> JuliaObj: return jl.Dense(self.lvl.create_jl_obj()) + @property def shape_type(self) -> tuple: return self.lvl.shape_type + (self.dim_type,) @@ -112,6 +114,7 @@ class SparseListFormat(NestedLevelFormat): def create_jl_obj(self) -> JuliaObj: return jl.SparseList(self.lvl.create_jl_obj()) + @property def shape_type(self) -> tuple: return self.lvl.shape_type + (self.dim_type,) @@ -133,6 +136,7 @@ class SparseCOOFormat(NestedLevelFormat): def create_jl_obj(self) -> JuliaObj: return jl.SparseCOO(self.lvl.create_jl_obj()) + @property def shape_type(self) -> tuple: if self.dim_type is None: return self.lvl.shape_type + (self.N * self.lvl.ndim,) @@ -154,6 +158,7 @@ class SparseByteMapFormat(NestedLevelFormat): def create_jl_obj(self) -> JuliaObj: return jl.SparseByteMap(self.lvl.create_jl_obj()) + @property def shape_type(self) -> tuple: return self.lvl.shape_type + (self.dim_type,) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 01dc0df..7bc2e4b 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -197,6 +197,12 @@ def asarray( if isinstance(obj, np.ndarray): if dtype is not None: obj = np.asarray(obj, dtype=jl_dtypes.jl_to_np_dtype[dtype]) + + # np.asfortranarray converts 0-D arrays into shape-(1,) arrays, + # so keep scalar inputs on the dedicated rank-0 construction path. + if obj.ndim == 0: + return full((), obj.item(), dtype=dtype) + if copy: obj = obj.copy() if np.isfortran(obj) else np.asfortranarray(obj) else: From 011962bf6ba91c271c089ff8d2d19ea193ad1a3b Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Mon, 30 Mar 2026 10:42:17 -0400 Subject: [PATCH 64/81] fix --- array-api-skips.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/array-api-skips.txt b/array-api-skips.txt index f628b7f..e48443f 100644 --- a/array-api-skips.txt +++ b/array-api-skips.txt @@ -6,6 +6,9 @@ array_api_tests/test_special_cases.py::test_unary[sign((x_i is -0 or x_i == +0)) array_api_tests/test_searching_functions.py::test_where # `test_solve` is not defined in Finch, hangs as xfail array_api_tests/test_linalg.py::test_solve +# infinite rewriting recursion +array_api_tests/test_constants.py::test_inf +array_api_tests/test_constants.py::test_nan # test_signatures From e781375b09e40a20628224d848606deb0e15b3d0 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Tue, 31 Mar 2026 13:53:24 -0400 Subject: [PATCH 65/81] fix --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b6a1e02..29070d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,4 +61,4 @@ section-order = [ [tool.mypy] ignore_missing_imports = true -exclude = ["tests/reference"] +exclude = ["tests/reference", "array_api_tests"] From 9f0d3e2b26332187ec19e2eed13ae111dfefc266 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Tue, 31 Mar 2026 13:53:53 -0400 Subject: [PATCH 66/81] prepping for mypy --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 29070d0..d99f6fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,4 +61,4 @@ section-order = [ [tool.mypy] ignore_missing_imports = true -exclude = ["tests/reference", "array_api_tests"] +exclude = ["tests/reference", "(^|/)array_api_tests"] From 28dc8fb7860cb2224177a63ee9e4b5274582d7fc Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Mon, 6 Apr 2026 13:55:11 -0400 Subject: [PATCH 67/81] fix: arange finally working --- src/finch/__init__.py | 10 ++++--- src/finch/levels.py | 14 ++++++---- src/finch/tensor.py | 65 +++++++++++++++++++++++++++++++++++-------- 3 files changed, 69 insertions(+), 20 deletions(-) diff --git a/src/finch/__init__.py b/src/finch/__init__.py index b79b7ea..299d22a 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -41,7 +41,6 @@ greater, greater_equal, hypot, - imag, isfinite, isinf, isnan, @@ -74,7 +73,6 @@ pow, power, prod, - real, reciprocal, reduce, remainder, @@ -132,10 +130,13 @@ empty_like, full, full_like, + imag, linspace, ones, ones_like, + real, reshape, + where, zeros, zeros_like, ) @@ -185,10 +186,10 @@ "cos", "cosh", "divide", + "e", "einop", "einsum", "elementwise", - "e", "empty", "empty_like", "equal", @@ -246,8 +247,8 @@ "multiply", "nan", "negative", - "nextafter", "newaxis", + "nextafter", "not_equal", "ones", "ones_like", @@ -288,6 +289,7 @@ "uint64", "var", "vecdot", + "where", "zeros", "zeros_like", ] diff --git a/src/finch/levels.py b/src/finch/levels.py index 9f9f72a..c73c62e 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -134,12 +134,17 @@ class SparseCOOFormat(NestedLevelFormat): dim_type: tuple | None = np.intp def create_jl_obj(self) -> JuliaObj: - return jl.SparseCOO(self.lvl.create_jl_obj()) + coo_type = jl.seval(f"Finch.SparseCOO{{{self.N}}}") + return coo_type(self.lvl.create_jl_obj()) + + @property + def ndim(self) -> np.intp: + return self.lvl.ndim + np.intp(self.N) # FIXME: not sure about this fix. @property def shape_type(self) -> tuple: if self.dim_type is None: - return self.lvl.shape_type + (self.N * self.lvl.ndim,) + return self.lvl.shape_type + (np.intp,) * self.N return self.lvl.shape_type + self.dim_type @@ -198,9 +203,8 @@ def jlobj_to_format(obj: JuliaObj) -> LevelFormat: if jl.isa(obj, jl.Finch.SparseList): return SparseListFormat(jlobj_to_format(obj.lvl), type(obj.shape)) if jl.isa(obj, jl.Finch.SparseCOO): - return SparseCOOFormat( - jlobj_to_format(obj.lvl), type(obj).__type_params__[0], type(obj.shape) - ) + N = jl.seval("Finch.level_ndims")(jl.typeof(obj)) + return SparseCOOFormat(jlobj_to_format(obj.lvl), N, None) if jl.isa(obj, jl.Finch.SparseByteMap): return SparseByteMapFormat(jlobj_to_format(obj.lvl), type(obj.shape)) raise Exception("Unhandled exception!") diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 7bc2e4b..d007a40 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -1,5 +1,5 @@ -from typing import Any import operator +from typing import Any import numpy as np @@ -111,7 +111,10 @@ def _is_dense(self) -> bool: def todense(self) -> np.ndarray: obj = self._obj - if self._is_dense: + if self.ndim == 0: # early return for 0-D tensor. + return np.array(jl.fill_value(obj)) + + if self._is_dense(): # don't materialize a dense finch tensor shape = jl.size(obj) dense_tensor = obj.lvl @@ -149,9 +152,14 @@ def __array_namespace__(self, *, api_version: str | None = None) -> Any: return finch + def copy(self) -> "FinchJLTensor": + return FinchJLTensor(jl.deepcopy(self._obj)) + def _scalar_value(self): if self.shape != (): - raise TypeError("only 0-dimensional arrays can be converted to Python scalars") + raise TypeError( + "only 0-dimensional arrays can be converted to Python scalars" + ) return jl.getindex(self._obj) def __bool__(self) -> bool: @@ -234,7 +242,9 @@ def asarray( ) m, n = obj.shape if dtype is not None: - fill_value = np.asarray(fill_value, dtype=jl_dtypes.jl_to_np_dtype[dtype]).item() + fill_value = np.asarray( + fill_value, dtype=jl_dtypes.jl_to_np_dtype[dtype] + ).item() if obj.format == "coo": return FinchJLTensor( jl.Tensor( @@ -317,7 +327,7 @@ def full_like( fill_value: jl_dtypes.number, *, dtype: DType | None = None, - format: str = "coo", + format=None, ) -> FinchJLTensor: return full(x.shape, fill_value, dtype=dtype, format=format) @@ -326,7 +336,7 @@ def ones( shape: int | tuple[int, ...], *, dtype: DType | None = None, - format: str = "coo", + format=None, ) -> FinchJLTensor: return full(shape, np.float64(1), dtype=dtype, format=format) @@ -336,7 +346,7 @@ def ones_like( /, *, dtype: DType | None = None, - format: str = "coo", + format=None, ) -> FinchJLTensor: dtype = x.dtype if dtype is None else dtype return ones(x.shape, dtype=dtype, format=format) @@ -346,7 +356,7 @@ def zeros( shape: int | tuple[int, ...], *, dtype: DType | None = None, - format: str = "coo", + format=None, ) -> FinchJLTensor: return full(shape, np.float64(0), dtype=dtype, format=format) @@ -356,7 +366,7 @@ def zeros_like( /, *, dtype: DType | None = None, - format: str = "coo", + format=None, ) -> FinchJLTensor: dtype = x.dtype if dtype is None else dtype return zeros(x.shape, dtype=dtype, format=format) @@ -366,7 +376,7 @@ def empty( shape: int | tuple[int, ...], *, dtype: DType | None = None, - format: str = "coo", + format=None, ) -> FinchJLTensor: return full(shape, np.float64(0), dtype=dtype, format=format) @@ -376,7 +386,7 @@ def empty_like( /, *, dtype: DType | None = None, - format: str = "coo", + format=None, ) -> FinchJLTensor: dtype = x.dtype if dtype is None else dtype return empty(x.shape, dtype=dtype, format=format) @@ -393,6 +403,39 @@ def arange( return asarray(np.arange(start, stop, step, jl_dtypes.jl_to_np_dtype[dtype])) +def real( # finchlite versions caused infinite recursion. + x: FinchJLTensor, + /, + *, + dtype: DType | None = None, +) -> FinchJLTensor: + return asarray(np.real(x.todense()), dtype=jl_dtypes.jl_to_np_dtype[dtype]) + + +def imag( # finchlite versions caused infinite recursion. + x: FinchJLTensor, + /, + *, + dtype: DType | None = None, +) -> FinchJLTensor: + return asarray(np.imag(x.todense()), dtype=jl_dtypes.jl_to_np_dtype[dtype]) + + +def _to_numpy(x): + if isinstance(x, FinchJLTensor): + return x.todense() + return np.asarray(x) + + +def where( + condition, + x1, + x2, + /, +) -> FinchJLTensor: + return asarray(np.where(_to_numpy(condition), _to_numpy(x1), _to_numpy(x2))) + + def linspace( start: complex, stop: complex, From 29c575c62242e2619fe0c64e57d72ec59fbd5762 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Tue, 23 Jun 2026 10:59:57 -0400 Subject: [PATCH 68/81] fix: finchlite fix --- array-api-skips.txt | 2 ++ pyproject.toml | 2 +- src/finch/tensor.py | 17 ++++++++++------- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/array-api-skips.txt b/array-api-skips.txt index e48443f..bc839f8 100644 --- a/array-api-skips.txt +++ b/array-api-skips.txt @@ -199,6 +199,7 @@ array_api_tests/test_has_names.py::test_has_names[manipulation-squeeze] array_api_tests/test_has_names.py::test_has_names[manipulation-stack] array_api_tests/test_has_names.py::test_has_names[manipulation-tile] array_api_tests/test_has_names.py::test_has_names[manipulation-unstack] +array_api_tests/test_has_names.py::test_has_names[manipulation-reshape] array_api_tests/test_has_names.py::test_has_names[sorting-argsort] array_api_tests/test_has_names.py::test_has_names[sorting-sort] array_api_tests/test_has_names.py::test_has_names[data_type-isdtype] @@ -284,6 +285,7 @@ array_api_tests/test_manipulation_functions.py::test_stack array_api_tests/test_manipulation_functions.py::test_unstack array_api_tests/test_manipulation_functions.py::test_repeat array_api_tests/test_manipulation_functions.py::test_tile +array_api_tests/test_manipulation_functions.py::test_reshape # test_searching_functions diff --git a/pyproject.toml b/pyproject.toml index 3bb6c60..fe24468 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ dependencies = [ "numpy (>=1.19,<2.4)", "juliacall (>=0.9.24,<0.10.0)", "lark (>=1.3.0,<2.0.0)", - "finch-tensor-lite (==0.3.0)", + "finch-tensor-lite (==0.5.0)", ] [tool.poetry] diff --git a/src/finch/tensor.py b/src/finch/tensor.py index d007a40..eda268f 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -190,12 +190,12 @@ def asarray( ) -> FinchJLTensor: if fill_value is None: fill_value = 0.0 - if copy is None: - copy = True if isinstance(obj, FinchJLTensor): if copy: return obj.copy() return obj + if copy is None: + copy = True if isinstance(obj, int | float | complex | bool | list): if copy is False: raise ValueError( @@ -226,10 +226,13 @@ def asarray( return FinchJLTensor(jl.Tensor(lvl)) if hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): if copy: - if obj.format in ("coo", "csr"): + if obj.format == "csr": obj = obj.copy() if obj.has_sorted_indices else obj.sorted_indices() if not obj.has_canonical_format: obj.sum_duplicates() + elif obj.format == "coo": + obj = obj.copy() + obj.sum_duplicates() else: obj = obj.asformat("csr") if ( @@ -250,11 +253,11 @@ def asarray( jl.Tensor( jl.SparseCOOLevel( (n, m), - jl.ElementLevel(dtype, fill_value, obj.data), + jl.ElementLevel(fill_value, obj.data), 2, idxs=( - jl.Finch.PlusOneVector(obj.cols), - jl.Finch.PlusOneVector(obj.rows), + jl.Finch.PlusOneVector(obj.col), + jl.Finch.PlusOneVector(obj.row), ), ) ) @@ -264,7 +267,7 @@ def asarray( jl.Tensor( jl.DenseLevel( jl.SparseListLevel( - jl.ElementLevel(dtype, fill_value, obj.data), + jl.ElementLevel(fill_value, obj.data), m, jl.Finch.PlusOneVector(obj.indptr), jl.Finch.PlusOneVector(obj.indices), From 8a0415e4404e8001fe17713713209ad35e3fa9ea Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Wed, 24 Jun 2026 13:28:15 -0400 Subject: [PATCH 69/81] fix: bump types to finchlite --- src/finch/__init__.py | 8 ++-- src/finch/dtypes.py | 89 ++++++++++++++++++++----------------------- src/finch/tensor.py | 6 +-- src/finch/typing.py | 3 +- tests/test_asarray.py | 13 ++++--- 5 files changed, 57 insertions(+), 62 deletions(-) diff --git a/src/finch/__init__.py b/src/finch/__init__.py index 299d22a..116c899 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -13,7 +13,7 @@ atan2, atanh, bitwise_and, - bitwise_inverse, + bitwise_invert, bitwise_left_shift, bitwise_or, bitwise_right_shift, @@ -37,7 +37,7 @@ expm1, flatten, floor, - floordiv, + floor_divide, greater, greater_equal, hypot, @@ -166,7 +166,7 @@ "atan2", "atanh", "bitwise_and", - "bitwise_inverse", + "bitwise_invert", "bitwise_left_shift", "bitwise_or", "bitwise_right_shift", @@ -202,7 +202,7 @@ "float32", "float64", "floor", - "floordiv", + "floor_divide", "full", "full_like", "fuse", diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py index 2b95252..1e08273 100644 --- a/src/finch/dtypes.py +++ b/src/finch/dtypes.py @@ -2,60 +2,53 @@ import numpy as np -from .julia import jl - -int_: jl.DataType = jl.Int -int8: jl.DataType = jl.Int8 -int16: jl.DataType = jl.Int16 -int32: jl.DataType = jl.Int32 -int64: jl.DataType = jl.Int64 -uint: jl.DataType = jl.UInt -uint8: jl.DataType = jl.UInt8 -uint16: jl.DataType = jl.UInt16 -uint32: jl.DataType = jl.UInt32 -uint64: jl.DataType = jl.UInt64 -float16: jl.DataType = jl.Float16 -float32: jl.DataType = jl.Float32 -float64: jl.DataType = jl.Float64 -complex64: jl.DataType = jl.ComplexF32 -complex128: jl.DataType = jl.ComplexF64 -bool: jl.DataType = jl.Bool - -number: jl.DataType = jl.Number -complex: jl.DataType = jl.Complex -integer: jl.DataType = jl.Integer -abstract_float: jl.DataType = jl.AbstractFloat +import finchlite as fl +from finchlite.algebra.ftypes import FType + +int8: FType = fl.int8 +int16: FType = fl.int16 +int32: FType = fl.int32 +int64: FType = fl.int64 +int_: FType = fl.intp +uint8: FType = fl.uint8 +uint16: FType = fl.uint16 +uint32: FType = fl.uint32 +uint64: FType = fl.uint64 +uint: FType = uint32 if np.uintp == np.uint32 else uint64 +float16: FType = fl.float16 +float32: FType = fl.float32 +float64: FType = fl.float64 +complex64: FType = fl.complex64 +complex128: FType = fl.complex128 +bool: FType = fl.bool + +number = builtins.int | builtins.float | builtins.complex | builtins.bool + +finfo = fl.finfo +iinfo = fl.iinfo jl_to_np_dtype = { - int_: np.int_, - int8: np.int8, - int16: np.int16, - int32: np.int32, - int64: np.int64, - uint: np.uint, - uint8: np.uint8, - uint16: np.uint16, - uint32: np.uint32, - uint64: np.uint64, - float16: np.float16, - float32: np.float32, - float64: np.float64, - complex64: np.complex64, - complex128: np.complex128, - bool: builtins.bool, + int_: int_.dtype, + int8: int8.dtype, + int16: int16.dtype, + int32: int32.dtype, + int64: int64.dtype, + uint: uint.dtype, + uint8: uint8.dtype, + uint16: uint16.dtype, + uint32: uint32.dtype, + uint64: uint64.dtype, + float16: float16.dtype, + float32: float32.dtype, + float64: float64.dtype, + complex64: complex64.dtype, + complex128: complex128.dtype, + bool: bool.dtype, None: None, } -def finfo(dtype): - return np.finfo(jl_to_np_dtype[dtype]) - - -def iinfo(dtype): - return np.iinfo(jl_to_np_dtype[dtype]) - - def can_cast(from_, to, /) -> builtins.bool: - if hasattr(from_, "dtype"): + if not isinstance(from_, FType) and hasattr(from_, "dtype"): from_ = from_.dtype return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index eda268f..3c897ec 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -13,7 +13,7 @@ SparseCOOFormat, jlobj_to_format, ) -from .typing import DType, JuliaObj +from .typing import DType, JuliaObj, number from .utils import add_missing_dims, add_plus_one, expand_ellipsis @@ -294,7 +294,7 @@ def reshape( def full( shape: int | tuple[int, ...], - val: jl_dtypes.number, + val: number, *, dtype: DType | None = None, format=None, @@ -327,7 +327,7 @@ def full( def full_like( x: FinchJLTensor, /, - fill_value: jl_dtypes.number, + fill_value: number, *, dtype: DType | None = None, format=None, diff --git a/src/finch/typing.py b/src/finch/typing.py index d3003ab..b154aec 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -1,5 +1,6 @@ import juliacall as jc +from finchlite.algebra.ftypes import FType JuliaObj = jc.AnyValue -DType = jc.AnyValue +DType = FType number = int | float | bool | complex diff --git a/tests/test_asarray.py b/tests/test_asarray.py index 5bc5998..267b4bc 100644 --- a/tests/test_asarray.py +++ b/tests/test_asarray.py @@ -4,6 +4,7 @@ import numpy as np +import finch from finch import asarray from finch.tensor import FinchJLTensor @@ -51,7 +52,7 @@ def test_asarray_with_fill_value(self): def test_asarray_with_dtype(self): """Test asarray with explicit dtype.""" arr = np.array([[1, 2], [3, 4]], dtype=np.int32) - result = asarray(arr, dtype=np.int32) + result = asarray(arr, dtype=finch.int32) assert isinstance(result, FinchJLTensor) def test_asarray_default_fill_value(self): @@ -313,9 +314,9 @@ def test_asarray_invalid_type(self): asarray("invalid string input") def test_asarray_invalid_list(self): - """Test asarray with plain Python list (should fail).""" + """Test asarray with copy=False on a plain Python list (should fail).""" with pytest.raises((ValueError, TypeError, AttributeError)): - asarray([1, 2, 3]) + asarray([1, 2, 3], copy=False) def test_asarray_dict_input(self): """Test asarray with dict input.""" @@ -340,11 +341,11 @@ def test_asarray_copy_none(self): def test_asarray_all_options(self): """Test asarray with all options specified.""" arr = np.array([[1.0, 2.0], [3.0, 4.0]]) - result = asarray(arr, dtype=np.float64, fill_value=0.0, copy=True) + result = asarray(arr, dtype=finch.float64, fill_value=0.0, copy=True) assert isinstance(result, FinchJLTensor) def test_asarray_numpy_no_copy(self): - """Test asarray on numpy with copy=False.""" - arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + """Test asarray on a Fortran-order numpy array with copy=False.""" + arr = np.asfortranarray([[1.0, 2.0], [3.0, 4.0]]) result = asarray(arr, copy=False) assert isinstance(result, FinchJLTensor) From 03468d26b5a1ab67d2ac2e8448504c34f431b978 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Wed, 24 Jun 2026 13:42:14 -0400 Subject: [PATCH 70/81] fix: version bump fixes --- src/finch/compiler.py | 8 ++++--- src/finch/tensor.py | 11 ++++++--- tests/test_compiler.py | 54 +++++++++++++++++++++++++++++------------- 3 files changed, 50 insertions(+), 23 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 920e941..b5ec928 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -2,7 +2,9 @@ import operator import finchlite.finch_notation.nodes as ntn -from finchlite.algebra import make_tuple, overwrite, promote_max, promote_min +from finchlite.algebra.ffuncs import make_tuple, overwrite +from finchlite.algebra.ffuncs import max as fl_max +from finchlite.algebra.ffuncs import min as fl_min from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary @@ -13,8 +15,8 @@ red_ops_map = { operator.add: "+", operator.mul: "*", - promote_max: "<>", - promote_min: "<>", + fl_max: "<>", + fl_min: "<>", } ops_to_ignore = [make_tuple] diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 3c897ec..151d028 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -3,7 +3,7 @@ import numpy as np -from finchlite import EagerTensor, Tensor, TensorFType +from finchlite import Tensor, TensorFType from . import dtypes as jl_dtypes from .julia import jc, jl @@ -42,7 +42,7 @@ def dtype(self) -> Any: def shape_type(self) -> tuple: return tuple(reversed(self._lvl.shape_type)) - def __call__(self, shape: tuple) -> Tensor: + def construct(self, shape: tuple) -> Tensor: return FinchJLTensor( jl.Finch.Tensor(self._lvl.create_jl_obj(), tuple(reversed(shape))) ) @@ -50,6 +50,11 @@ def __call__(self, shape: tuple) -> Tensor: def from_numpy(self, _) -> Tensor: raise NotImplementedError + def __call__(self, val: Any) -> Tensor: + raise NotImplementedError( + f"Tensor conversion not yet implemented for {type(self).__name__}" + ) + def __eq__(self, other): if not isinstance(other, FinchJLTensorFType): return False @@ -59,7 +64,7 @@ def __hash__(self): return hash(("FinchJLTensorFType", self._lvl)) -class FinchJLTensor(EagerTensor): +class FinchJLTensor(Tensor): def __init__(self, obj: JuliaObj): if isinstance(obj, JuliaObj): assert jl.isa(obj, jl.Finch.Tensor) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 1fde3ad..2844bd6 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -27,6 +27,7 @@ Variable, ) +import finch from finch.compiler import FinchJLCompiler, FinchJLKernel from finch.julia import jl from finch.levels import DenseFormat, ElementFormat @@ -51,21 +52,27 @@ Block( ( Assign( - Variable("m", ExtentFType(np.int64, np.int64)), + Variable( + "m", ExtentFType(finch.int64, finch.int64) + ), Call( Literal(dimension), (Variable("A", a_format), Literal(0)), ), ), Assign( - Variable("n", ExtentFType(np.int64, np.int64)), + Variable( + "n", ExtentFType(finch.int64, finch.int64) + ), Call( Literal(dimension), (Variable("B", a_format), Literal(1)), ), ), Assign( - Variable("p", ExtentFType(np.int64, np.int64)), + Variable( + "p", ExtentFType(finch.int64, finch.int64) + ), Call( Literal(dimension), (Variable("A", a_format), Literal(1)), @@ -79,20 +86,29 @@ Literal(0.0), Literal(operator.add), ( - Variable("m", ExtentFType(np.int64, np.int64)), - Variable("n", ExtentFType(np.int64, np.int64)), + Variable( + "m", ExtentFType(finch.int64, finch.int64) + ), + Variable( + "n", ExtentFType(finch.int64, finch.int64) + ), ), ), Loop( - Variable("i", np.int64), - Variable("m", ExtentFType(np.int64, np.int64)), + Variable("i", finch.int64), + Variable( + "m", ExtentFType(finch.int64, finch.int64) + ), Loop( - Variable("k", np.int64), - Variable("p", ExtentFType(np.int64, np.int64)), + Variable("k", finch.int64), + Variable( + "p", ExtentFType(finch.int64, finch.int64) + ), Loop( - Variable("j", np.int64), + Variable("j", finch.int64), Variable( - "n", ExtentFType(np.int64, np.int64) + "n", + ExtentFType(finch.int64, finch.int64), ), Block( ( @@ -103,8 +119,12 @@ Literal(operator.add) ), ( - Variable("i", np.int64), - Variable("j", np.int64), + Variable( + "i", finch.int64 + ), + Variable( + "j", finch.int64 + ), ), ), Call( @@ -120,11 +140,11 @@ ( Variable( "i", - np.int64, + finch.int64, ), Variable( "k", - np.int64, + finch.int64, ), ), ) @@ -139,11 +159,11 @@ ( Variable( "k", - np.int64, + finch.int64, ), Variable( "j", - np.int64, + finch.int64, ), ), ) From caac7eb5eebf8c6ef8aa7f3a85407efa20b85c78 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Wed, 24 Jun 2026 14:12:26 -0400 Subject: [PATCH 71/81] fix: version bump fixes --- src/finch/compiler.py | 15 ++++++--------- src/finch/dtypes.py | 33 +++++++++++++++++++++++++++++++++ src/finch/levels.py | 27 ++++++++++++++++++--------- tests/test_compiler.py | 13 +++++-------- 4 files changed, 62 insertions(+), 26 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index b5ec928..5375edb 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,22 +1,19 @@ import math -import operator import finchlite.finch_notation.nodes as ntn -from finchlite.algebra.ffuncs import make_tuple, overwrite -from finchlite.algebra.ffuncs import max as fl_max -from finchlite.algebra.ffuncs import min as fl_min +from finchlite.algebra.ffuncs import add, eq, make_tuple, max, min, mul, overwrite from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary from .julia import jl from .tensor import FinchJLTensor -ops_map = {operator.add: "+", operator.mul: "*", operator.eq: "=="} +ops_map = {add: "+", mul: "*", eq: "=="} red_ops_map = { - operator.add: "+", - operator.mul: "*", - fl_max: "<>", - fl_min: "<>", + add: "+", + mul: "*", + max: "<>", + min: "<>", } ops_to_ignore = [make_tuple] diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py index 1e08273..d4d8c17 100644 --- a/src/finch/dtypes.py +++ b/src/finch/dtypes.py @@ -5,6 +5,8 @@ import finchlite as fl from finchlite.algebra.ftypes import FType +from .julia import jl + int8: FType = fl.int8 int16: FType = fl.int16 int32: FType = fl.int32 @@ -52,3 +54,34 @@ def can_cast(from_, to, /) -> builtins.bool: if not isinstance(from_, FType) and hasattr(from_, "dtype"): from_ = from_.dtype return np.can_cast(jl_to_np_dtype[from_], jl_to_np_dtype[to]) + + +# Julia DataType -> finchlite FType, used when reading dtypes back out of +# raw Julia Finch tensor objects (see levels.jlobj_to_format). +jl_dtype_to_fl = { + jl.Int8: int8, + jl.Int16: int16, + jl.Int32: int32, + jl.Int64: int64, + jl.UInt8: uint8, + jl.UInt16: uint16, + jl.UInt32: uint32, + jl.UInt64: uint64, + jl.Float16: float16, + jl.Float32: float32, + jl.Float64: float64, + jl.ComplexF32: complex64, + jl.ComplexF64: complex128, + jl.Bool: bool, +} + + +def to_fl_dtype(x) -> FType: + """Normalize a Julia DataType, numpy dtype/scalar type, Python builtin + type, or finchlite FType into the corresponding finchlite FType.""" + if isinstance(x, FType): + return x + fl_dtype = jl_dtype_to_fl.get(x) + if fl_dtype is not None: + return fl_dtype + return fl.ftype(x) diff --git a/src/finch/levels.py b/src/finch/levels.py index c73c62e..9b90069 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -4,6 +4,7 @@ import numpy as np +from . import dtypes from .julia import jl from .typing import JuliaObj, number @@ -47,7 +48,9 @@ class ElementFormat(LevelFormat): def __init__(self, fill_value: number, element_type: Any | None = None): self._fill_value = fill_value - self._element_type = type(fill_value) if element_type is None else element_type + self._element_type = dtypes.to_fl_dtype( + type(fill_value) if element_type is None else element_type + ) @property def ndim(self) -> np.intp: @@ -89,7 +92,7 @@ class DenseFormat(NestedLevelFormat): """ lvl: NestedLevelFormat - dim_type: type = np.intp + dim_type: Any = dtypes.int_ def create_jl_obj(self) -> JuliaObj: return jl.Dense(self.lvl.create_jl_obj()) @@ -109,7 +112,7 @@ class SparseListFormat(NestedLevelFormat): """ lvl: NestedLevelFormat - dim_type: type = np.intp + dim_type: Any = dtypes.int_ def create_jl_obj(self) -> JuliaObj: return jl.SparseList(self.lvl.create_jl_obj()) @@ -131,7 +134,7 @@ class SparseCOOFormat(NestedLevelFormat): lvl: NestedLevelFormat N: int = 2 - dim_type: tuple | None = np.intp + dim_type: tuple | None = dtypes.int_ def create_jl_obj(self) -> JuliaObj: coo_type = jl.seval(f"Finch.SparseCOO{{{self.N}}}") @@ -158,7 +161,7 @@ class SparseByteMapFormat(NestedLevelFormat): """ lvl: NestedLevelFormat - dim_type: type = np.intp + dim_type: Any = dtypes.int_ def create_jl_obj(self) -> JuliaObj: return jl.SparseByteMap(self.lvl.create_jl_obj()) @@ -196,15 +199,21 @@ def jlobj_to_format(obj: JuliaObj) -> LevelFormat: obj_type = jl.typeof(obj) return ElementFormat( jl.Finch.level_fill_value(obj_type), - jl.Finch.level_eltype(obj_type), + dtypes.to_fl_dtype(jl.Finch.level_eltype(obj_type)), ) if jl.isa(obj, jl.Finch.Dense): - return DenseFormat(jlobj_to_format(obj.lvl), type(obj.shape)) + return DenseFormat( + jlobj_to_format(obj.lvl), dtypes.to_fl_dtype(type(obj.shape)) + ) if jl.isa(obj, jl.Finch.SparseList): - return SparseListFormat(jlobj_to_format(obj.lvl), type(obj.shape)) + return SparseListFormat( + jlobj_to_format(obj.lvl), dtypes.to_fl_dtype(type(obj.shape)) + ) if jl.isa(obj, jl.Finch.SparseCOO): N = jl.seval("Finch.level_ndims")(jl.typeof(obj)) return SparseCOOFormat(jlobj_to_format(obj.lvl), N, None) if jl.isa(obj, jl.Finch.SparseByteMap): - return SparseByteMapFormat(jlobj_to_format(obj.lvl), type(obj.shape)) + return SparseByteMapFormat( + jlobj_to_format(obj.lvl), dtypes.to_fl_dtype(type(obj.shape)) + ) raise Exception("Unhandled exception!") diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 2844bd6..9ac6f52 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -1,9 +1,8 @@ -import operator - import pytest import numpy as np +from finchlite.algebra.ffuncs import add, mul from finchlite.compile import ExtentFType, dimension from finchlite.finch_notation.nodes import ( Access, @@ -84,7 +83,7 @@ Declare( Slot("C_", a_format), Literal(0.0), - Literal(operator.add), + Literal(add), ( Variable( "m", ExtentFType(finch.int64, finch.int64) @@ -115,9 +114,7 @@ Increment( Access( Slot("C_", a_format), - Update( - Literal(operator.add) - ), + Update(Literal(add)), ( Variable( "i", finch.int64 @@ -128,7 +125,7 @@ ), ), Call( - Literal(operator.mul), + Literal(mul), ( Unwrap( Access( @@ -176,7 +173,7 @@ ), ), ), - Freeze(Slot("C_", a_format), Literal(operator.add)), + Freeze(Slot("C_", a_format), Literal(add)), Repack(Slot("C_", a_format), Variable("C", a_format)), Return(Variable("C", a_format)), ), From 3b65f3e4b2af2032e062c7d1508c6ec799984311 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Wed, 24 Jun 2026 15:16:00 -0400 Subject: [PATCH 72/81] column major fixes --- src/finch/dtypes.py | 6 ++++ src/finch/tensor.py | 82 +++++++++++++++++++++++++++++++++++---------- 2 files changed, 70 insertions(+), 18 deletions(-) diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py index d4d8c17..3c27554 100644 --- a/src/finch/dtypes.py +++ b/src/finch/dtypes.py @@ -85,3 +85,9 @@ def to_fl_dtype(x) -> FType: if fl_dtype is not None: return fl_dtype return fl.ftype(x) + + +# finchlite FType -> Julia DataType, the inverse of jl_dtype_to_fl, used when +# a real Julia type is needed (e.g. juliacall.convert) for an FType obtained +# from a tensor's dtype/element_type. +fl_dtype_to_jl = {v: k for k, v in jl_dtype_to_fl.items()} diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 151d028..77391cf 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -43,9 +43,12 @@ def shape_type(self) -> tuple: return tuple(reversed(self._lvl.shape_type)) def construct(self, shape: tuple) -> Tensor: - return FinchJLTensor( - jl.Finch.Tensor(self._lvl.create_jl_obj(), tuple(reversed(shape))) - ) + # `jl.Finch.Tensor(lvl, dims)` sets the resulting tensor's raw + # `jl.size` directly to `dims` (no implicit row/col-major swap, unlike + # construction from real array data), and `FinchJLTensor.shape` reports + # that raw size unreversed -- so `shape` must be passed through as-is + # to keep it in the same (Python) axis order the caller gave us. + return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), shape)) def from_numpy(self, _) -> Tensor: raise NotImplementedError @@ -99,11 +102,36 @@ def __getitem__(self, key): key = add_missing_dims(key, self.shape) key = add_plus_one(key, self.shape) - result = jl.getindex(self._obj, *reversed(key)) if all(isinstance(k, int) for k in key): - return result - assert jl.isa(result, jl.Finch.Tensor) - return FinchJLTensor(result) + return jl.getindex(self._obj, *key) + + # Finch's getindex has no notion of `None`/newaxis, so strip those + # entries out, index normally, then re-insert size-1 axes into the + # result at the positions implied by the original key. + real_key = tuple(k for k in key if k is not jl.nothing) + newaxis_positions = [] + axis = 0 + for k in key: + if k is jl.nothing: + newaxis_positions.append(axis) + axis += 1 + elif not isinstance(k, int): + axis += 1 + + result = jl.getindex(self._obj, *real_key) + + if not newaxis_positions: + assert jl.isa(result, jl.Finch.Tensor) + return FinchJLTensor(result) + + arr = ( + FinchJLTensor(result).todense() + if jl.isa(result, jl.Finch.Tensor) + else np.asarray(result) + ) + for pos in newaxis_positions: + arr = np.expand_dims(arr, pos) + return asarray(arr) def _is_dense(self) -> bool: lvl = self._obj.lvl @@ -126,7 +154,9 @@ def todense(self) -> np.ndarray: else: # create materialized dense array shape = jl.size(obj) - dense_lvls = jl.Element(jc.convert(self.dtype, jl.fill_value(obj))) + dense_lvls = jl.Element( + jc.convert(jl_dtypes.fl_dtype_to_jl[self.dtype], jl.fill_value(obj)) + ) for _ in range(self.ndim): dense_lvls = jl.Dense(dense_lvls) dense_tensor = jl.Tensor(dense_lvls, obj).lvl # materialize @@ -254,15 +284,33 @@ def asarray( fill_value, dtype=jl_dtypes.jl_to_np_dtype[dtype] ).item() if obj.format == "coo": + # SparseCOOLevel keeps shape/axis order as given (no row<->col + # swap, unlike Dense levels: COO has no flat buffer to + # reinterpret with reversed strides, so there's nothing to + # transpose there). + # Finch's COO reader does not sort or dedupe its inputs, but does + # expect them pre-sorted in column-major scan order (row fastest, + # col slowest); np.lexsort sorts by its *last* key first, so this + # sorts by col primarily and row secondarily. + order = np.lexsort((obj.row, obj.col)) + row_s = obj.row[order] + col_s = obj.col[order] + data_s = obj.data[order] + nnz = len(data_s) return FinchJLTensor( jl.Tensor( jl.SparseCOOLevel( - (n, m), - jl.ElementLevel(fill_value, obj.data), - 2, - idxs=( - jl.Finch.PlusOneVector(obj.col), - jl.Finch.PlusOneVector(obj.row), + jl.ElementLevel(fill_value, data_s), + (m, n), + # ptr marks the single coordinate block [1, nnz+1); + # it's a plain Python list, so it needs an explicit + # jl.Vector to become a real Julia array (numpy arrays + # get this automatically via PythonCall's zero-copy + # PyArray wrapping). + jl.Vector([1, nnz + 1]), + ( + jl.Finch.PlusOneVector(row_s), + jl.Finch.PlusOneVector(col_s), ), ) ) @@ -323,10 +371,8 @@ def full( format = SparseCOOFormat(ElementFormat(val, dtype), len(shape)) if format.fill_value != val: - return FinchJLTensor( - jl.Tensor(format.create_jl_obj(), np.full(val, reversed(shape))) - ) - return FinchJLTensor(jl.Tensor(format.create_jl_obj(), *reversed(shape))) + return FinchJLTensor(jl.Tensor(format.create_jl_obj(), np.full(shape, val))) + return FinchJLTensor(jl.Tensor(format.create_jl_obj(), *shape)) def full_like( From 06c758a812f5e15d53ad5d29612e5fda0291adb2 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Wed, 24 Jun 2026 17:05:14 -0400 Subject: [PATCH 73/81] fix: update tests to use asarray --- src/finch/__init__.py | 7 ++++ src/finch/compiler.py | 2 +- src/finch/tensor.py | 74 +++++++++++++++++++++++++----------------- tests/test_einsum.py | 38 ++++++++++------------ tests/test_indexing.py | 7 ++-- 5 files changed, 72 insertions(+), 56 deletions(-) diff --git a/src/finch/__init__.py b/src/finch/__init__.py index 116c899..a88e6e3 100644 --- a/src/finch/__init__.py +++ b/src/finch/__init__.py @@ -38,6 +38,7 @@ flatten, floor, floor_divide, + get_default_scheduler, greater, greater_equal, hypot, @@ -77,6 +78,7 @@ reduce, remainder, round, + set_default_scheduler, sign, signbit, sin, @@ -141,6 +143,11 @@ zeros_like, ) +# finch's whole purpose is being the Julia-backed array API implementation, +# so it should default to executing through Finch.jl rather than finchlite's +# generic fallback interpreter. +set_default_scheduler(COMPILE_JULIA) + e = math.e pi = math.pi inf = math.inf diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 5375edb..d9f9846 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -124,7 +124,7 @@ def generate_julia(self, prgm, nestingLvl=0): case ntn.Access(tns, _, idxs): tns_str = self.generate_julia(tns, nestingLvl) idx_str = ",".join( - [self.generate_julia(idx, nestingLvl) for idx in idxs] + [self.generate_julia(idx, nestingLvl) for idx in reversed(idxs)] ) return f"{tns_str}[{idx_str}]" diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 77391cf..3c2da2e 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -43,12 +43,12 @@ def shape_type(self) -> tuple: return tuple(reversed(self._lvl.shape_type)) def construct(self, shape: tuple) -> Tensor: - # `jl.Finch.Tensor(lvl, dims)` sets the resulting tensor's raw - # `jl.size` directly to `dims` (no implicit row/col-major swap, unlike - # construction from real array data), and `FinchJLTensor.shape` reports - # that raw size unreversed -- so `shape` must be passed through as-is - # to keep it in the same (Python) axis order the caller gave us. - return FinchJLTensor(jl.Finch.Tensor(self._lvl.create_jl_obj(), shape)) + # EXPERIMENTAL reversed-axis convention: jl.size is always kept as + # the reverse of the Python-facing shape; FinchJLTensor.shape un- + # reverses it back on the way out (see there for the full rationale). + return FinchJLTensor( + jl.Finch.Tensor(self._lvl.create_jl_obj(), tuple(reversed(shape))) + ) def from_numpy(self, _) -> Tensor: raise NotImplementedError @@ -86,8 +86,13 @@ def dtype(self) -> Any: @property def shape(self) -> tuple: - """Shape of the tensor.""" - return jl.size(self._obj) + """Shape of the tensor. + + EXPERIMENTAL: jl.size is always kept as the reverse of the + Python-facing shape (see asarray/full/__getitem__/todense), so this + un-reverses it back to Python axis order. + """ + return tuple(reversed(jl.size(self._obj))) def __getitem__(self, key): if not isinstance(key, tuple): @@ -103,7 +108,7 @@ def __getitem__(self, key): key = add_plus_one(key, self.shape) if all(isinstance(k, int) for k in key): - return jl.getindex(self._obj, *key) + return jl.getindex(self._obj, *reversed(key)) # Finch's getindex has no notion of `None`/newaxis, so strip those # entries out, index normally, then re-insert size-1 axes into the @@ -118,7 +123,7 @@ def __getitem__(self, key): elif not isinstance(k, int): axis += 1 - result = jl.getindex(self._obj, *real_key) + result = jl.getindex(self._obj, *reversed(real_key)) if not newaxis_positions: assert jl.isa(result, jl.Finch.Tensor) @@ -164,8 +169,11 @@ def todense(self) -> np.ndarray: for _ in range(self.ndim): dense_tensor = dense_tensor.lvl + # `shape` here is jl.size(obj), i.e. the reversed-axis shape; reshape + # into that, then transpose (a cheap stride-only view) back to the + # Python-facing axis order. arr = jl.reshape(dense_tensor.val, tuple(shape)) - return np.asarray(arr) + return np.asarray(arr).transpose() def __eq__(self, other): return isinstance(other, FinchJLTensor) and self._obj == other._obj @@ -241,22 +249,25 @@ def asarray( if dtype is not None: obj = np.asarray(obj, dtype=jl_dtypes.jl_to_np_dtype[dtype]) - # np.asfortranarray converts 0-D arrays into shape-(1,) arrays, - # so keep scalar inputs on the dedicated rank-0 construction path. if obj.ndim == 0: return full((), obj.item(), dtype=dtype) + # EXPERIMENTAL reversed-axis convention: keep the buffer in its + # natural C (row-major) layout -- which is exactly the column-major + # layout of the *reversed*-shape tensor -- instead of permuting data + # to Fortran order, and build the DenseLevel nest in reverse axis + # order so jl.size ends up as reversed(obj.shape). if copy: - obj = obj.copy() if np.isfortran(obj) else np.asfortranarray(obj) + obj = obj.copy() if obj.flags["C_CONTIGUOUS"] else np.ascontiguousarray(obj) else: - if not np.isfortran(obj): + if not obj.flags["C_CONTIGUOUS"]: raise ValueError( "Unable to avoid copy while creating an array as requested." ) - buf = np.reshape(np.permute_dims(obj, tuple(reversed(range(obj.ndim)))), -1) + buf = np.reshape(obj, -1) lvl = jl.ElementLevel(np.asarray(fill_value, dtype=obj.dtype).item(), buf) - for i in obj.shape: + for i in reversed(obj.shape): lvl = jl.DenseLevel(lvl, i) return FinchJLTensor(jl.Tensor(lvl)) if hasattr(obj, "__module__") and obj.__module__.startswith("scipy.sparse"): @@ -284,15 +295,12 @@ def asarray( fill_value, dtype=jl_dtypes.jl_to_np_dtype[dtype] ).item() if obj.format == "coo": - # SparseCOOLevel keeps shape/axis order as given (no row<->col - # swap, unlike Dense levels: COO has no flat buffer to - # reinterpret with reversed strides, so there's nothing to - # transpose there). - # Finch's COO reader does not sort or dedupe its inputs, but does - # expect them pre-sorted in column-major scan order (row fastest, - # col slowest); np.lexsort sorts by its *last* key first, so this - # sorts by col primarily and row secondarily. - order = np.lexsort((obj.row, obj.col)) + # EXPERIMENTAL reversed-axis convention: shape and coordinate + # order are both given reversed ((n, m) and (col, row)), so axis + # 0 of the SparseCOOLevel is "col". Sorting must vary axis 0 + # (col) fastest -- i.e. row primary, col secondary -- so col is + # lexsort's secondary (first) key and row its primary (last) key. + order = np.lexsort((obj.col, obj.row)) row_s = obj.row[order] col_s = obj.col[order] data_s = obj.data[order] @@ -301,7 +309,7 @@ def asarray( jl.Tensor( jl.SparseCOOLevel( jl.ElementLevel(fill_value, data_s), - (m, n), + (n, m), # ptr marks the single coordinate block [1, nnz+1); # it's a plain Python list, so it needs an explicit # jl.Vector to become a real Julia array (numpy arrays @@ -309,8 +317,8 @@ def asarray( # PyArray wrapping). jl.Vector([1, nnz + 1]), ( - jl.Finch.PlusOneVector(row_s), jl.Finch.PlusOneVector(col_s), + jl.Finch.PlusOneVector(row_s), ), ) ) @@ -371,8 +379,14 @@ def full( format = SparseCOOFormat(ElementFormat(val, dtype), len(shape)) if format.fill_value != val: - return FinchJLTensor(jl.Tensor(format.create_jl_obj(), np.full(shape, val))) - return FinchJLTensor(jl.Tensor(format.create_jl_obj(), *shape)) + # jl.Tensor(lvl, real_array) infers jl.size literally from the + # array's own .shape (no implicit reversal on that path -- proven + # separately), so under the reversed-axis convention the array + # passed here must itself already be shaped in reverse. + return FinchJLTensor( + jl.Tensor(format.create_jl_obj(), np.full(tuple(reversed(shape)), val)) + ) + return FinchJLTensor(jl.Tensor(format.create_jl_obj(), *reversed(shape))) def full_like( diff --git a/tests/test_einsum.py b/tests/test_einsum.py index 32f92f0..aeab9cb 100644 --- a/tests/test_einsum.py +++ b/tests/test_einsum.py @@ -1,9 +1,9 @@ import numpy as np import finchlite -from juliacall import Main as jl -from finch import COMPILE_JULIA, FinchJLTensor +import finch +from finch import COMPILE_JULIA def test_pass_through(rng): @@ -11,7 +11,7 @@ def test_pass_through(rng): A = rng.random((5, 5)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + A_finch = finch.asarray(A) B = finchlite.einop("B[i,j] = A[i,j]", A=A_finch) np.allclose(B.todense(), A) @@ -22,7 +22,7 @@ def test_transpose(rng): A = rng.random((5, 5)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + A_finch = finch.asarray(A) B = finchlite.einop("B[i,j] = A[j, i]", A=A_finch) np.allclose(B.todense(), A.T) @@ -34,8 +34,8 @@ def test_basic_addition_with_transpose(rng): B = rng.random((5, 5)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) - B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), B)) + A_finch = finch.asarray(A) + B_finch = finch.asarray(B) C = finchlite.einop("C[i,j] = A[i,j] + B[j,i]", A=A_finch, B=B_finch) C_ref = A + B.T @@ -48,8 +48,8 @@ def test_matrix_multiplication(rng): B = rng.random((4, 5)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) - B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), B)) + A_finch = finch.asarray(A) + B_finch = finch.asarray(B) C = finchlite.einop("C[i,j] += A[i,k] * B[k,j]", A=A_finch, B=B_finch) C_ref = A @ B @@ -62,8 +62,8 @@ def test_element_wise_multiplication(rng): B = rng.random((4, 4)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) - B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), B)) + A_finch = finch.asarray(A) + B_finch = finch.asarray(B) C = finchlite.einop("C[i,j] = A[i,j] * B[i,j]", A=A_finch, B=B_finch) C_ref = A * B @@ -75,7 +75,7 @@ def test_sum_reduction(rng): A = rng.random((3, 4)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + A_finch = finch.asarray(A) C = finchlite.einop("C[i] += A[i,j]", A=A_finch) C_ref = np.sum(A, axis=1) @@ -87,7 +87,7 @@ def test_maximum_reduction(rng): A = rng.random((3, 4)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + A_finch = finch.asarray(A) C = finchlite.einop("C[i] max= A[i,j]", A=A_finch) C_ref = np.max(A, axis=1) @@ -100,8 +100,8 @@ def test_outer_product(rng): B = rng.random(4) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Element(0.0)), A)) - B_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Element(0.0)), B)) + A_finch = finch.asarray(A) + B_finch = finch.asarray(B) C = finchlite.einop("C[i,j] = A[i] * B[j]", A=A_finch, B=B_finch) C_ref = np.outer(A, B) @@ -114,12 +114,8 @@ def test_batch_matrix_multiplication(rng): B = rng.random((2, 4, 5)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor( - jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Dense(jl.Element(0.0)))), A) - ) - B_finch = FinchJLTensor( - jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Dense(jl.Element(0.0)))), B) - ) + A_finch = finch.asarray(A) + B_finch = finch.asarray(B) C = finchlite.einop("C[b,i,j] += A[b,i,k] * B[b,k,j]", A=A_finch, B=B_finch) C_ref = np.matmul(A, B) @@ -131,7 +127,7 @@ def test_minimum_reduction(rng): A = rng.random((3, 4)) finchlite.interface.set_default_scheduler(ctx=COMPILE_JULIA) - A_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0.0))), A)) + A_finch = finch.asarray(A) C = finchlite.einop("C[i] min= A[i,j]", A=A_finch) C_ref = np.min(A, axis=1) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 5bd2c03..626dbd6 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -4,6 +4,7 @@ from juliacall import Main as jl +import finch from finch import FinchJLTensor @@ -49,7 +50,7 @@ def test_indexing_1d(arr1d, index): ], ) def test_indexing_2d(arr2d, index): - arr_finch = FinchJLTensor(jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Element(0))), arr2d)) + arr_finch = finch.asarray(arr2d) actual = arr_finch[index] expected = arr2d[index] @@ -88,9 +89,7 @@ def test_indexing_2d(arr2d, index): ], ) def test_indexing_3d(arr3d, index): - arr_finch = FinchJLTensor( - jl.Finch.Tensor(jl.Dense(jl.Dense(jl.Dense(jl.Element(0)))), arr3d) - ) + arr_finch = finch.asarray(arr3d) actual = arr_finch[index] expected = arr3d[index] From 2793e9d1c6d6c3514cdf612e62b8948c62ce59de Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Wed, 24 Jun 2026 17:09:03 -0400 Subject: [PATCH 74/81] fix: idk --- tests/test_asarray.py | 9 +++++++-- tests/test_compiler.py | 4 +++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/test_asarray.py b/tests/test_asarray.py index 267b4bc..7d46451 100644 --- a/tests/test_asarray.py +++ b/tests/test_asarray.py @@ -345,7 +345,12 @@ def test_asarray_all_options(self): assert isinstance(result, FinchJLTensor) def test_asarray_numpy_no_copy(self): - """Test asarray on a Fortran-order numpy array with copy=False.""" - arr = np.asfortranarray([[1.0, 2.0], [3.0, 4.0]]) + """Test asarray on a C-order numpy array with copy=False. + + Under the reversed-axis storage convention, the buffer is kept in + its natural C (row-major) layout (see asarray's Dense branch), so + C-contiguous -- not Fortran-order -- is what allows a zero copy. + """ + arr = np.ascontiguousarray([[1.0, 2.0], [3.0, 4.0]]) result = asarray(arr, copy=False) assert isinstance(result, FinchJLTensor) diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 9ac6f52..3d1d18f 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -181,13 +181,15 @@ ), ) ), + # Access index order is reversed in codegen to match the + # reversed-axis storage convention (see compiler.py). """function matmul(C,A,B) @finch C .= 0.0 @finch begin for i = _ for k = _ for j = _ - C[i,j] += *(A[i,k],B[k,j]) + C[j,i] += *(A[k,i],B[j,k]) end end end From db9a3b48e2872d3f9273905eec26dcdb753654d7 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 25 Jun 2026 09:33:43 -0400 Subject: [PATCH 75/81] fix --- src/finch/compiler.py | 35 +++++++++++++++++++++++++++++++++-- src/finch/levels.py | 13 +++++++++++-- 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index d9f9846..d2d1c3d 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,14 +1,40 @@ import math +import numpy as np + import finchlite.finch_notation.nodes as ntn -from finchlite.algebra.ffuncs import add, eq, make_tuple, max, min, mul, overwrite +from finchlite.algebra.ffuncs import ( + add, + eq, + equal, + greater, + greater_equal, + less, + less_equal, + make_tuple, + max, + min, + mul, + not_equal, + overwrite, +) from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary from .julia import jl from .tensor import FinchJLTensor -ops_map = {add: "+", mul: "*", eq: "=="} +ops_map = { + add: "+", + mul: "*", + eq: "==", + equal: "==", + not_equal: "!=", + less: "<", + less_equal: "<=", + greater: ">", + greater_equal: ">=", +} red_ops_map = { add: "+", mul: "*", @@ -190,6 +216,11 @@ def generate_julia(self, prgm, nestingLvl=0): return self.pack_dict[name] case ntn.Literal(val): + # Julia booleans are lowercase, unlike Python's str(bool). + # numpy.bool_ is not a subclass of Python's bool, so check + # both. + if isinstance(val, bool | np.bool_): + return "true" if val else "false" # Julia represents inf differently than how its represented in python if val > 0 and math.isinf(val): return "Inf" diff --git a/src/finch/levels.py b/src/finch/levels.py index 9b90069..bb89d87 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -75,7 +75,16 @@ def __hash__(self): return hash((self.__class__.__name__, self._fill_value, self._element_type)) def create_jl_obj(self) -> JuliaObj: - return jl.Element(self._fill_value) + # Cast through element_type so the real Julia element type matches + # the requested dtype, rather than whatever Julia infers from + # self._fill_value's own Python/numpy type. + val = self.fill_value + # PythonCall wraps numpy scalars (e.g. numpy.bool_) as 0-d PyArrays + # rather than native Julia values, and Finch's ElementLevel rejects + # non-isbits defaults -- unwrap to a native Python scalar first. + if isinstance(val, np.generic): + val = val.item() + return jl.Element(val) @property def shape_type(self) -> tuple: @@ -147,7 +156,7 @@ def ndim(self) -> np.intp: @property def shape_type(self) -> tuple: if self.dim_type is None: - return self.lvl.shape_type + (np.intp,) * self.N + return self.lvl.shape_type + (dtypes.int_,) * self.N return self.lvl.shape_type + self.dim_type From 8c42e1d37b7142f1b83c261426b6e2c9f7fab6b6 Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 25 Jun 2026 10:37:49 -0400 Subject: [PATCH 76/81] fix: update ops map --- src/finch/compiler.py | 117 +++++++++++++++++++++++++++++++----------- 1 file changed, 87 insertions(+), 30 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index d2d1c3d..eb25796 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -2,44 +2,80 @@ import numpy as np +import finchlite.algebra.ffuncs as ffuncs import finchlite.finch_notation.nodes as ntn -from finchlite.algebra.ffuncs import ( - add, - eq, - equal, - greater, - greater_equal, - less, - less_equal, - make_tuple, - max, - min, - mul, - not_equal, - overwrite, -) +from finchlite.algebra.algebra import FinchOperator +from finchlite.algebra.ffuncs import add, make_tuple, max, min, mul, overwrite from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary from .julia import jl from .tensor import FinchJLTensor -ops_map = { - add: "+", - mul: "*", - eq: "==", - equal: "==", - not_equal: "!=", - less: "<", - less_equal: "<=", - greater: ">", - greater_equal: ">=", +_JULIA_NAME_OVERRIDES = { + "add": "+", + "mul": "*", + "sub": "-", + "truediv": "/", + "divide": "/", + "floordiv": "div", + "mod": "mod", + "remainder": "mod", + "pow": "^", + "neg": "-", + "pos": "+", + "eq": "==", + "equal": "==", + "ne": "!=", + "not_equal": "!=", + "lt": "<", + "less": "<", + "le": "<=", + "less_equal": "<=", + "gt": ">", + "greater": ">", + "ge": ">=", + "greater_equal": ">=", + "and_": "&", + "or_": "|", + "not_": "!", + "invert": "~", + "lshift": "<<", + "rshift": ">>", + "divmod": "divrem", + "logical_and": "&", + "logical_or": "|", + "logical_not": "!", + "logical_xor": "xor", + "square": "abs2", + "reciprocal": "inv", + "atan2": "atan", + "conjugate": "conj", + "where": "ifelse", + "clip": "clamp", + "truth": "Bool", } + + +def _build_ops_map() -> dict: + m = {} + for py_name in dir(ffuncs): + obj = getattr(ffuncs, py_name) + if isinstance(obj, FinchOperator): + m[obj] = _JULIA_NAME_OVERRIDES.get(py_name, py_name) + return m + + +ops_map = _build_ops_map() red_ops_map = { add: "+", mul: "*", max: "<>", min: "<>", + ffuncs.and_: "&", + ffuncs.or_: "|", + ffuncs.logical_and: "&", + ffuncs.logical_or: "|", } ops_to_ignore = [make_tuple] @@ -53,7 +89,13 @@ def __init__(self, func_name, jl_code): def __call__(self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...]: finch_fn = getattr(jl, self.func_name) - result = finch_fn(*[arg._obj for arg in args]) + # Some args may be finchlite's plain-Python Scalar (rank-0, not + # backed by a Julia object) rather than a FinchJLTensor -- pass its + # raw value through directly. + raw_args = [ + arg._obj if isinstance(arg, FinchJLTensor) else arg.val for arg in args + ] + result = finch_fn(*raw_args) # The finch function returns tuples when multiple values are returned # or a non-tuple when a single value is returned. @@ -109,10 +151,18 @@ def generate_julia(self, prgm, nestingLvl=0): return "" tab_str = " " * nestingLvl - return ( - f"{tab_str}{self.generate_julia(lhs, nestingLvl)} = " + stmt = ( + f"{self.generate_julia(lhs, nestingLvl)} = " f"{self.generate_julia(rhs, nestingLvl)}" ) + # A rank-0 (loop-less) Assign falls outside any enclosing + # @finch block, so Tensor access would otherwise go through + # Finch's generic (non-macro) setindex!/getindex, which has + # real bugs for some element types (e.g. Bool). Wrap it the + # same way Declare already wraps its single-line form. + if not self.in_finch_block: + return f"{tab_str}@finch {stmt}" + return f"{tab_str}{stmt}" case ntn.Declare(tns, init, op, _): # TODO: what is the purpose of op here @@ -190,8 +240,15 @@ def generate_julia(self, prgm, nestingLvl=0): # If the operation is overwrite just codegen an assignment if lhs.mode.op.val == overwrite: - return f"{tab_str}{lhs_str} = {rhs_str}" - return f"{tab_str}{lhs_str} {red_ops_map[lhs.mode.op.val]}= {rhs_str}" + stmt = f"{lhs_str} = {rhs_str}" + else: + stmt = f"{lhs_str} {red_ops_map[lhs.mode.op.val]}= {rhs_str}" + # A rank-0 (loop-less) Increment falls outside any enclosing + # @finch block -- see the matching comment in the Assign + # case above for why that needs wrapping. + if not self.in_finch_block: + return f"{tab_str}@finch {stmt}" + return f"{tab_str}{stmt}" case ntn.Unwrap(arg): return self.generate_julia(arg, nestingLvl) From 9eade271a1f77443885e800d9da885157b6530ad Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 25 Jun 2026 11:55:06 -0400 Subject: [PATCH 77/81] fix: array creation functions pass --- src/finch/compiler.py | 21 ++++++++++++++++++--- src/finch/tensor.py | 19 +++++++++++++++++-- 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index eb25796..4340734 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -12,6 +12,17 @@ from .julia import jl from .tensor import FinchJLTensor + +def _wrap_scalar(val): + # A bare jl.Element(val)/raw Julia scalar isn't a valid Finch tensor + # access root inside @finch -- build a real rank-0 tensor with an + # explicit length-1 buffer (see the matching fix in tensor.py's full()). + if isinstance(val, np.generic): + val = val.item() + buf = np.asarray([val]) + return jl.Tensor(jl.ElementLevel(buf.item(), buf)) + + _JULIA_NAME_OVERRIDES = { "add": "+", "mul": "*", @@ -90,10 +101,14 @@ def __init__(self, func_name, jl_code): def __call__(self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...]: finch_fn = getattr(jl, self.func_name) # Some args may be finchlite's plain-Python Scalar (rank-0, not - # backed by a Julia object) rather than a FinchJLTensor -- pass its - # raw value through directly. + # backed by a Julia object) rather than a FinchJLTensor. The + # generated code accesses every argument with Finch's tensor-access + # syntax (tns[]) inside an @finch block, which requires a real + # wrapped Finch tensor object -- a bare Julia scalar isn't a valid + # access root. So wrap the raw value as a minimal rank-0 tensor. raw_args = [ - arg._obj if isinstance(arg, FinchJLTensor) else arg.val for arg in args + arg._obj if isinstance(arg, FinchJLTensor) else _wrap_scalar(arg.val) + for arg in args ] result = finch_fn(*raw_args) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 3c2da2e..41e6c86 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -182,6 +182,10 @@ def __repr__(self): return jl.sprint(jl.show, self._obj) def __str__(self): + # A 0-d tensor has no axes to permute, and an empty swizzle + # permutation crashes Finch's SwizzleArray size/show machinery. + if self.ndim == 0: + return jl.sprint(jl.show, jl.MIME("text/plain"), self._obj) swiz = jl.swizzle(self._obj, tuple(reversed(range(self.ndim, 1, -1)))) return jl.sprint(jl.show, jl.MIME("text/plain"), swiz) @@ -373,7 +377,15 @@ def full( # Rank-0 tensors should be represented as a leaf element level. # Building them through SparseCOO requires an explicit rank parameter. if len(shape) == 0 and format is None: - return FinchJLTensor(jl.Tensor(ElementFormat(val, dtype).create_jl_obj())) + elt_fmt = ElementFormat(val, dtype) + cast_val = elt_fmt.fill_value + if isinstance(cast_val, np.generic): + cast_val = cast_val.item() + # A bare jl.Element(val) treats `val` as the level's default and + # allocates an empty buffer, so reading it back gives the type-zero, + # not `val`. A rank-0 tensor needs an explicit length-1 buffer. + buf = np.asarray([cast_val], dtype=dtype) + return FinchJLTensor(jl.Tensor(jl.ElementLevel(cast_val, buf))) if format is None: format = SparseCOOFormat(ElementFormat(val, dtype), len(shape)) @@ -384,7 +396,10 @@ def full( # separately), so under the reversed-axis convention the array # passed here must itself already be shaped in reverse. return FinchJLTensor( - jl.Tensor(format.create_jl_obj(), np.full(tuple(reversed(shape)), val)) + jl.Tensor( + format.create_jl_obj(), + np.full(tuple(reversed(shape)), val, dtype=dtype), + ) ) return FinchJLTensor(jl.Tensor(format.create_jl_obj(), *reversed(shape))) From 793e2f38a63876b27193183334d603e2f41cbf7d Mon Sep 17 00:00:00 2001 From: sreevickrant Date: Thu, 25 Jun 2026 14:09:52 -0400 Subject: [PATCH 78/81] fix: FinchJLTensor inherits from OverrideTensor --- src/finch/compiler.py | 125 +++++++++++++++++++----------------------- src/finch/levels.py | 9 +-- src/finch/tensor.py | 24 +++++++- 3 files changed, 83 insertions(+), 75 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 4340734..68a9e30 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -1,11 +1,9 @@ -import math - import numpy as np import finchlite.algebra.ffuncs as ffuncs import finchlite.finch_notation.nodes as ntn from finchlite.algebra.algebra import FinchOperator -from finchlite.algebra.ffuncs import add, make_tuple, max, min, mul, overwrite +from finchlite.algebra.ffuncs import make_tuple, overwrite from finchlite.compile import NotationCompiler, dimension from finchlite.finch_assembly import AssemblyKernel, AssemblyLibrary @@ -13,17 +11,21 @@ from .tensor import FinchJLTensor -def _wrap_scalar(val): - # A bare jl.Element(val)/raw Julia scalar isn't a valid Finch tensor - # access root inside @finch -- build a real rank-0 tensor with an - # explicit length-1 buffer (see the matching fix in tensor.py's full()). +def _scalar_to_jl(val): + # Fallback for finchlite's plain Scalar (raw Python values that never + # went through our tensor system). Build a real rank-0 Finch tensor so + # the generated @finch code can access it via tns[] syntax. if isinstance(val, np.generic): val = val.item() buf = np.asarray([val]) return jl.Tensor(jl.ElementLevel(buf.item(), buf)) -_JULIA_NAME_OVERRIDES = { +# Single source of truth for Python ffunc name → Julia operator/function name. +# max/min use Finch's <> semiring syntax; they appear only as Aggregate +# (reduction) nodes, never as plain element-wise Calls, so this is safe. +_JULIA_NAMES = { + # arithmetic "add": "+", "mul": "*", "sub": "-", @@ -35,6 +37,7 @@ def _wrap_scalar(val): "pow": "^", "neg": "-", "pos": "+", + # comparisons "eq": "==", "equal": "==", "ne": "!=", @@ -47,17 +50,22 @@ def _wrap_scalar(val): "greater": ">", "ge": ">=", "greater_equal": ">=", + # bitwise / logical "and_": "&", "or_": "|", "not_": "!", "invert": "~", "lshift": "<<", "rshift": ">>", - "divmod": "divrem", "logical_and": "&", "logical_or": "|", "logical_not": "!", "logical_xor": "xor", + # reductions (Finch <>= semiring syntax) + "max": "<>", + "min": "<>", + # misc + "divmod": "divrem", "square": "abs2", "reciprocal": "inv", "atan2": "atan", @@ -67,27 +75,30 @@ def _wrap_scalar(val): "truth": "Bool", } - -def _build_ops_map() -> dict: - m = {} - for py_name in dir(ffuncs): - obj = getattr(ffuncs, py_name) - if isinstance(obj, FinchOperator): - m[obj] = _JULIA_NAME_OVERRIDES.get(py_name, py_name) - return m - - -ops_map = _build_ops_map() -red_ops_map = { - add: "+", - mul: "*", - max: "<>", - min: "<>", - ffuncs.and_: "&", - ffuncs.or_: "|", - ffuncs.logical_and: "&", - ffuncs.logical_or: "|", +# Names of ffuncs that are valid as reduction operators. +_REDUCTION_OPS = { + "add", + "mul", + "max", + "min", + "and_", + "or_", + "logical_and", + "logical_or", } + + +def _ops_for(names=None) -> dict: + """Build a FinchOperator → Julia-name map from _JULIA_NAMES.""" + return { + obj: _JULIA_NAMES.get(n, n) + for n in (names if names is not None else dir(ffuncs)) + if isinstance(obj := getattr(ffuncs, n, None), FinchOperator) + } + + +ops_map = _ops_for() +red_ops_map = _ops_for(_REDUCTION_OPS) ops_to_ignore = [make_tuple] @@ -100,14 +111,8 @@ def __init__(self, func_name, jl_code): def __call__(self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ...]: finch_fn = getattr(jl, self.func_name) - # Some args may be finchlite's plain-Python Scalar (rank-0, not - # backed by a Julia object) rather than a FinchJLTensor. The - # generated code accesses every argument with Finch's tensor-access - # syntax (tns[]) inside an @finch block, which requires a real - # wrapped Finch tensor object -- a bare Julia scalar isn't a valid - # access root. So wrap the raw value as a minimal rank-0 tensor. raw_args = [ - arg._obj if isinstance(arg, FinchJLTensor) else _wrap_scalar(arg.val) + arg._obj if isinstance(arg, FinchJLTensor) else _scalar_to_jl(arg.val) for arg in args ] result = finch_fn(*raw_args) @@ -157,9 +162,7 @@ def generate_julia(self, prgm, nestingLvl=0): return body_str case ntn.Assign(lhs, rhs): - # TODO: Can we make this better? - # Special condition to ignore all assigns associated with - # finding loop bounds + # Ignore assigns used only to find loop bounds. if isinstance(rhs, ntn.Dimension) or ( isinstance(rhs, ntn.Call) and rhs.op.val == dimension ): @@ -179,8 +182,7 @@ def generate_julia(self, prgm, nestingLvl=0): return f"{tab_str}@finch {stmt}" return f"{tab_str}{stmt}" - case ntn.Declare(tns, init, op, _): - # TODO: what is the purpose of op here + case ntn.Declare(tns, init, _, _): tab_str = " " * nestingLvl return ( f"{tab_str}@finch {self.generate_julia(tns, nestingLvl)} .= " @@ -193,20 +195,17 @@ def generate_julia(self, prgm, nestingLvl=0): case ntn.Loop(idx, _, body): tab_str = " " * nestingLvl - tab_str_1 = " " * (nestingLvl + 1) idx_str = self.generate_julia(idx, nestingLvl) - - is_outermost_loop = False - if self.in_finch_block is False: - is_outermost_loop = True + outermost = not self.in_finch_block + if outermost: self.in_finch_block = True - loop_body = self.generate_julia(body, nestingLvl + 2) - else: - loop_body = self.generate_julia(body, nestingLvl + 1) - - if not is_outermost_loop: + loop_body = self.generate_julia( + body, nestingLvl + (2 if outermost else 1) + ) + if not outermost: return f"{tab_str}for {idx_str} = _\n{loop_body}{tab_str}end\n" self.in_finch_block = False + tab_str_1 = " " * (nestingLvl + 1) return ( f"{tab_str}@finch begin\n{tab_str_1}for {idx_str} = " f"_\n{loop_body}{tab_str_1}end\n{tab_str}end" @@ -234,6 +233,7 @@ def generate_julia(self, prgm, nestingLvl=0): return f"{tab_str}if {cond_str}\n{body_str}\n{tab_str}end" case ntn.IfElse(cond, then_body, else_body): + tab_str = " " * nestingLvl cond_str = self.generate_julia(cond, nestingLvl) then_body_str = self.generate_julia(then_body, nestingLvl + 1) else_body_str = self.generate_julia(else_body, nestingLvl + 1) @@ -246,14 +246,6 @@ def generate_julia(self, prgm, nestingLvl=0): tab_str = " " * nestingLvl lhs_str = self.generate_julia(lhs, nestingLvl) rhs_str = self.generate_julia(rhs, nestingLvl) - - # TODO: Is this the correct assumption to make - if not ( - isinstance(lhs, ntn.Access) and isinstance(lhs.mode, ntn.Update) - ): - raise Exception("Increment expects the lhs to be an access") - - # If the operation is overwrite just codegen an assignment if lhs.mode.op.val == overwrite: stmt = f"{lhs_str} = {rhs_str}" else: @@ -269,7 +261,6 @@ def generate_julia(self, prgm, nestingLvl=0): return self.generate_julia(arg, nestingLvl) case ntn.Unpack(lhs, rhs): - # TODO: Is this the right assumption to make if not isinstance(rhs, ntn.Variable): raise Exception("The unpack was not called with variable as RHS.") self.pack_dict[lhs.name] = self.generate_julia(rhs, nestingLvl) @@ -288,21 +279,15 @@ def generate_julia(self, prgm, nestingLvl=0): return self.pack_dict[name] case ntn.Literal(val): - # Julia booleans are lowercase, unlike Python's str(bool). - # numpy.bool_ is not a subclass of Python's bool, so check - # both. + # Julia booleans are lowercase; numpy.bool_ is not a bool subclass. if isinstance(val, bool | np.bool_): return "true" if val else "false" - # Julia represents inf differently than how its represented in python - if val > 0 and math.isinf(val): - return "Inf" - if val < 0 and math.isinf(val): - return "-Inf" + if isinstance(val, float | np.floating) and np.isinf(val): + return "Inf" if val > 0 else "-Inf" return str(val) case ntn.Variable(name, _): - # finch tensor lite uses character(#) in the naming of variables - # that however is not valid julia syntax + # finchlite uses '#' in generated names; not valid Julia syntax. return name.replace("#", "_") # TODO: Cached, Dimension, Thaw, Stack, Value are unimplemented. diff --git a/src/finch/levels.py b/src/finch/levels.py index bb89d87..3e5ad79 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -79,10 +79,11 @@ def create_jl_obj(self) -> JuliaObj: # the requested dtype, rather than whatever Julia infers from # self._fill_value's own Python/numpy type. val = self.fill_value - # PythonCall wraps numpy scalars (e.g. numpy.bool_) as 0-d PyArrays - # rather than native Julia values, and Finch's ElementLevel rejects - # non-isbits defaults -- unwrap to a native Python scalar first. - if isinstance(val, np.generic): + # PythonCall wraps numpy.bool_ as a 0-d PyArray (non-isbits), which + # Finch's ElementLevel rejects -- unwrap to a native Python bool. + # For all other numpy scalar types, pass them through directly so + # Julia preserves the right type (e.g. np.uint8(0) → UInt8, not Int64). + if isinstance(val, np.bool_): val = val.item() return jl.Element(val) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index 41e6c86..e1e43b0 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -4,6 +4,7 @@ import numpy as np from finchlite import Tensor, TensorFType +from finchlite.tensor.override_tensor import OverrideTensor from . import dtypes as jl_dtypes from .julia import jc, jl @@ -67,7 +68,28 @@ def __hash__(self): return hash(("FinchJLTensorFType", self._lvl)) -class FinchJLTensor(Tensor): +class FinchJLTensor(OverrideTensor): + def override_module(self): + import finch + + return finch + + def __array_function__(self, func, types, args, kwargs): + # Guard np.asarray specifically: lazy.asarray() calls np.asarray() + # internally, and redirecting that back to finch.asarray() creates an + # infinite loop (finch.asarray returns FinchJLTensor unchanged, then + # lazy.asarray calls np.asarray again, etc.). Returning NotImplemented + # lets numpy fall back to the dtype=object path in lazy.asarray, which + # correctly returns the tensor as-is for downstream Julia compilation. + import finch + + if func.__name__ == "asarray": + return NotImplemented + override_func = getattr(finch, func.__name__, None) + if override_func is None: + return NotImplemented + return override_func(*args, **kwargs) + def __init__(self, obj: JuliaObj): if isinstance(obj, JuliaObj): assert jl.isa(obj, jl.Finch.Tensor) From 48010dc30b531367bf422b84494da63fc78536c9 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 26 Jun 2026 13:04:04 -0400 Subject: [PATCH 79/81] should work *crosses fingers* --- src/finch/compiler.py | 61 ++++++++++++++++++------------------------ src/finch/dtypes.py | 7 +++++ src/finch/levels.py | 2 +- src/finch/tensor.py | 8 ++++-- src/finch/typing.py | 7 +++++ tests/test_compiler.py | 4 +-- 6 files changed, 49 insertions(+), 40 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 68a9e30..563cd7f 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -137,29 +137,30 @@ def __getattr__(self, name: str) -> FinchJLKernel: class FinchJLGenerator: def __init__(self): self.pack_dict = {} - self.in_finch_block = False def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: self.pack_dict.clear() - self.in_finch_block = False return self.generate_julia(prgm) def generate_julia(self, prgm, nestingLvl=0): match prgm: case ntn.Function(name, args, body): - body_str = self.generate_julia(body, nestingLvl + 1) - arg_str = ",".join( - [self.generate_julia(arg, nestingLvl) for arg in args] - ) - return f"function {name}({arg_str})\n{body_str}end" + body_str = self.generate_julia(body, nestingLvl + 2) + arg_strs = [] + for arg in args: + match(arg): + case ntn.Variable(sym, type): + arg_strs.append(f"{sym}") #TODO later use finch_kernel and type the args + case _: + raise NotImplementedError + arg_str = ",".join(arg_strs) + return f"function {name}({arg_str})\n @finch begin\n{body_str}\n end\nend" case ntn.Block(bodies): body_str = "" - for body in bodies: - curr_body_str = self.generate_julia(body, nestingLvl) - if curr_body_str != "": - body_str += f"{curr_body_str}\n" - return body_str + body_strs = [self.generate_julia(body, nestingLvl) for body in bodies] + body_strs = [body_str for body_str in body_strs if body_str != ""] + return "\n".join(body_strs) case ntn.Assign(lhs, rhs): # Ignore assigns used only to find loop bounds. @@ -173,19 +174,12 @@ def generate_julia(self, prgm, nestingLvl=0): f"{self.generate_julia(lhs, nestingLvl)} = " f"{self.generate_julia(rhs, nestingLvl)}" ) - # A rank-0 (loop-less) Assign falls outside any enclosing - # @finch block, so Tensor access would otherwise go through - # Finch's generic (non-macro) setindex!/getindex, which has - # real bugs for some element types (e.g. Bool). Wrap it the - # same way Declare already wraps its single-line form. - if not self.in_finch_block: - return f"{tab_str}@finch {stmt}" return f"{tab_str}{stmt}" case ntn.Declare(tns, init, _, _): tab_str = " " * nestingLvl return ( - f"{tab_str}@finch {self.generate_julia(tns, nestingLvl)} .= " + f"{tab_str}{self.generate_julia(tns, nestingLvl)} .= " f"{self.generate_julia(init, nestingLvl)}" ) @@ -196,19 +190,13 @@ def generate_julia(self, prgm, nestingLvl=0): case ntn.Loop(idx, _, body): tab_str = " " * nestingLvl idx_str = self.generate_julia(idx, nestingLvl) - outermost = not self.in_finch_block - if outermost: - self.in_finch_block = True loop_body = self.generate_julia( - body, nestingLvl + (2 if outermost else 1) + body, nestingLvl + 1 ) - if not outermost: - return f"{tab_str}for {idx_str} = _\n{loop_body}{tab_str}end\n" - self.in_finch_block = False tab_str_1 = " " * (nestingLvl + 1) return ( - f"{tab_str}@finch begin\n{tab_str_1}for {idx_str} = " - f"_\n{loop_body}{tab_str_1}end\n{tab_str}end" + f"{tab_str}for {idx_str} = _\n" + f"{loop_body}\n{tab_str}end" ) case ntn.Access(tns, _, idxs): @@ -250,11 +238,6 @@ def generate_julia(self, prgm, nestingLvl=0): stmt = f"{lhs_str} = {rhs_str}" else: stmt = f"{lhs_str} {red_ops_map[lhs.mode.op.val]}= {rhs_str}" - # A rank-0 (loop-less) Increment falls outside any enclosing - # @finch block -- see the matching comment in the Assign - # case above for why that needs wrapping. - if not self.in_finch_block: - return f"{tab_str}@finch {stmt}" return f"{tab_str}{stmt}" case ntn.Unwrap(arg): @@ -273,6 +256,12 @@ def generate_julia(self, prgm, nestingLvl=0): case ntn.Freeze(_, _): return "" + case ntn.Thaw(_, _): + return "" + + case ntn.Cached(_, _): + return "" + case ntn.Slot(name): if name not in self.pack_dict: raise Exception(f"{name} Slot does not exist in registry.") @@ -290,8 +279,8 @@ def generate_julia(self, prgm, nestingLvl=0): # finchlite uses '#' in generated names; not valid Julia syntax. return name.replace("#", "_") - # TODO: Cached, Dimension, Thaw, Stack, Value are unimplemented. case _: + # Dimension, Stack, Value are deliberately unimplemented. raise Exception(f"Unhandled node type: {type(prgm)}") @@ -302,6 +291,8 @@ def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: kernel_dict = {} for func in prgm.children: generated_prgm = generator(func) + print("-"*80) + print(generated_prgm) kernel_dict[func.name.name] = FinchJLKernel(func.name.name, generated_prgm) return FinchJLLibrary(kernel_dict) diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py index 3c27554..135a58b 100644 --- a/src/finch/dtypes.py +++ b/src/finch/dtypes.py @@ -4,6 +4,7 @@ import finchlite as fl from finchlite.algebra.ftypes import FType +from .typing import JLFType from .julia import jl @@ -91,3 +92,9 @@ def to_fl_dtype(x) -> FType: # a real Julia type is needed (e.g. juliacall.convert) for an FType obtained # from a tensor's dtype/element_type. fl_dtype_to_jl = {v: k for k, v in jl_dtype_to_fl.items()} + +def to_jl_type(T:FType): + if T in fl_dtype_to_jl: + return fl_dtype_to_jl[T] + elif isinstance(T, JLFType): + return T.to_jl_type() \ No newline at end of file diff --git a/src/finch/levels.py b/src/finch/levels.py index 3e5ad79..440d8d1 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -51,7 +51,7 @@ def __init__(self, fill_value: number, element_type: Any | None = None): self._element_type = dtypes.to_fl_dtype( type(fill_value) if element_type is None else element_type ) - + @property def ndim(self) -> np.intp: return np.intp(0) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index e1e43b0..bc928db 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -14,12 +14,12 @@ SparseCOOFormat, jlobj_to_format, ) -from .typing import DType, JuliaObj, number +from .typing import DType, JuliaObj, number, JLFType from .utils import add_missing_dims, add_plus_one, expand_ellipsis # Tensor Class and associated ftype -class FinchJLTensorFType(TensorFType): +class FinchJLTensorFType(TensorFType, JLFType): def __init__(self, lvl): self._lvl: LevelFormat = lvl @@ -42,6 +42,10 @@ def dtype(self) -> Any: @property def shape_type(self) -> tuple: return tuple(reversed(self._lvl.shape_type)) + + @property + def jl_type(self): + return jl.Finch.Tensor[self.format.jl_type] def construct(self, shape: tuple) -> Tensor: # EXPERIMENTAL reversed-axis convention: jl.size is always kept as diff --git a/src/finch/typing.py b/src/finch/typing.py index b154aec..f7a8403 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -1,6 +1,13 @@ import juliacall as jc from finchlite.algebra.ftypes import FType +from abc import ABC, abstractmethod JuliaObj = jc.AnyValue DType = FType number = int | float | bool | complex + +class JLFType(FType, ABC): + @property + @abstractmethod + def jl_type(self): + pass diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 3d1d18f..8a9438c 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -184,8 +184,8 @@ # Access index order is reversed in codegen to match the # reversed-axis storage convention (see compiler.py). """function matmul(C,A,B) - @finch C .= 0.0 @finch begin + C .= 0.0 for i = _ for k = _ for j = _ @@ -193,8 +193,8 @@ end end end + return C end - return C end""", ) ], From ba0983a8a67190e350622ffc9885cd673206e02c Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 26 Jun 2026 13:06:16 -0400 Subject: [PATCH 80/81] fix --- src/finch/compiler.py | 23 +++++++++++------------ src/finch/dtypes.py | 10 ++++++---- src/finch/levels.py | 2 +- src/finch/tensor.py | 4 ++-- src/finch/typing.py | 4 +++- src/finch/utils.py | 1 - 6 files changed, 23 insertions(+), 21 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index 563cd7f..f2ba676 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -148,13 +148,18 @@ def generate_julia(self, prgm, nestingLvl=0): body_str = self.generate_julia(body, nestingLvl + 2) arg_strs = [] for arg in args: - match(arg): + match arg: case ntn.Variable(sym, type): - arg_strs.append(f"{sym}") #TODO later use finch_kernel and type the args + arg_strs.append( + f"{sym}" + ) # TODO later use finch_kernel and type the args case _: raise NotImplementedError arg_str = ",".join(arg_strs) - return f"function {name}({arg_str})\n @finch begin\n{body_str}\n end\nend" + return ( + f"function {name}({arg_str})\n @finch begin\n" + f"{body_str}\n end\nend" + ) case ntn.Block(bodies): body_str = "" @@ -190,14 +195,8 @@ def generate_julia(self, prgm, nestingLvl=0): case ntn.Loop(idx, _, body): tab_str = " " * nestingLvl idx_str = self.generate_julia(idx, nestingLvl) - loop_body = self.generate_julia( - body, nestingLvl + 1 - ) - tab_str_1 = " " * (nestingLvl + 1) - return ( - f"{tab_str}for {idx_str} = _\n" - f"{loop_body}\n{tab_str}end" - ) + loop_body = self.generate_julia(body, nestingLvl + 1) + return f"{tab_str}for {idx_str} = _\n{loop_body}\n{tab_str}end" case ntn.Access(tns, _, idxs): tns_str = self.generate_julia(tns, nestingLvl) @@ -291,7 +290,7 @@ def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: kernel_dict = {} for func in prgm.children: generated_prgm = generator(func) - print("-"*80) + print("-" * 80) print(generated_prgm) kernel_dict[func.name.name] = FinchJLKernel(func.name.name, generated_prgm) diff --git a/src/finch/dtypes.py b/src/finch/dtypes.py index 135a58b..1e7caf7 100644 --- a/src/finch/dtypes.py +++ b/src/finch/dtypes.py @@ -4,9 +4,9 @@ import finchlite as fl from finchlite.algebra.ftypes import FType -from .typing import JLFType from .julia import jl +from .typing import JLFType int8: FType = fl.int8 int16: FType = fl.int16 @@ -93,8 +93,10 @@ def to_fl_dtype(x) -> FType: # from a tensor's dtype/element_type. fl_dtype_to_jl = {v: k for k, v in jl_dtype_to_fl.items()} -def to_jl_type(T:FType): + +def to_jl_type(T: FType): if T in fl_dtype_to_jl: return fl_dtype_to_jl[T] - elif isinstance(T, JLFType): - return T.to_jl_type() \ No newline at end of file + if isinstance(T, JLFType): + return T.to_jl_type() + raise NotImplementedError diff --git a/src/finch/levels.py b/src/finch/levels.py index 440d8d1..3e5ad79 100644 --- a/src/finch/levels.py +++ b/src/finch/levels.py @@ -51,7 +51,7 @@ def __init__(self, fill_value: number, element_type: Any | None = None): self._element_type = dtypes.to_fl_dtype( type(fill_value) if element_type is None else element_type ) - + @property def ndim(self) -> np.intp: return np.intp(0) diff --git a/src/finch/tensor.py b/src/finch/tensor.py index bc928db..412045b 100644 --- a/src/finch/tensor.py +++ b/src/finch/tensor.py @@ -14,7 +14,7 @@ SparseCOOFormat, jlobj_to_format, ) -from .typing import DType, JuliaObj, number, JLFType +from .typing import DType, JLFType, JuliaObj, number from .utils import add_missing_dims, add_plus_one, expand_ellipsis @@ -42,7 +42,7 @@ def dtype(self) -> Any: @property def shape_type(self) -> tuple: return tuple(reversed(self._lvl.shape_type)) - + @property def jl_type(self): return jl.Finch.Tensor[self.format.jl_type] diff --git a/src/finch/typing.py b/src/finch/typing.py index f7a8403..7fc598c 100644 --- a/src/finch/typing.py +++ b/src/finch/typing.py @@ -1,11 +1,13 @@ +from abc import ABC, abstractmethod + import juliacall as jc from finchlite.algebra.ftypes import FType -from abc import ABC, abstractmethod JuliaObj = jc.AnyValue DType = FType number = int | float | bool | complex + class JLFType(FType, ABC): @property @abstractmethod diff --git a/src/finch/utils.py b/src/finch/utils.py index 3afa309..16a7bf8 100644 --- a/src/finch/utils.py +++ b/src/finch/utils.py @@ -53,7 +53,6 @@ def _slice_plus_one(s: slice, size: int) -> range: else: stop = stop_default - return jl.range(start=start, step=step, stop=stop) From b7cb4f434141bc93297abb2bcaad3db4f4310462 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Fri, 26 Jun 2026 13:20:05 -0400 Subject: [PATCH 81/81] fix --- src/finch/compiler.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/finch/compiler.py b/src/finch/compiler.py index f2ba676..564cf48 100644 --- a/src/finch/compiler.py +++ b/src/finch/compiler.py @@ -119,8 +119,8 @@ def __call__(self, *args: tuple[FinchJLTensor, ...]) -> tuple[FinchJLTensor, ... # The finch function returns tuples when multiple values are returned # or a non-tuple when a single value is returned. - if not isinstance(result, tuple): - result = (result,) + if jl.isa(result, jl.Finch.Tensor): + return (FinchJLTensor(result),) return tuple(FinchJLTensor(res) for res in result) @@ -151,7 +151,7 @@ def generate_julia(self, prgm, nestingLvl=0): match arg: case ntn.Variable(sym, type): arg_strs.append( - f"{sym}" + f"{sym.replace('#', '_')}" ) # TODO later use finch_kernel and type the args case _: raise NotImplementedError @@ -290,8 +290,6 @@ def __call__(self, prgm: ntn.Module) -> FinchJLLibrary: kernel_dict = {} for func in prgm.children: generated_prgm = generator(func) - print("-" * 80) - print(generated_prgm) kernel_dict[func.name.name] = FinchJLKernel(func.name.name, generated_prgm) return FinchJLLibrary(kernel_dict)