finch_notation

Access dataclass

Bases: NotationTree, NotationExpression

Notation AST expression representing the value of tensor tns at the indices idx....

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Access(NotationTree, NotationExpression):
    """
    Notation AST expression representing the value of tensor `tns` at the indices
    `idx...`.
    """

    tns: NotationExpression
    mode: AccessMode
    idxs: tuple[NotationExpression, ...]

    @property
    def result_type(self):
        if len(self.idxs) == 0:
            return self.tns.result_type
        return AccessFType(self.tns.result_type)

    @classmethod
    def from_children(cls, tns, mode, *idxs):
        return cls(tns, mode, idxs)

    @property
    def children(self):
        return [self.tns, self.mode, *self.idxs]

AccessMode dataclass

Bases: NotationNode

Notation AST node representing the access mode of a tensor.

Source code in src/finch/finch_notation/nodes.py

class AccessMode(NotationNode):
    """
    Notation AST node representing the access mode of a tensor.
    """

Assign dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that defines lhs as having the value rhs.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Assign(NotationTree, NotationStatement):
    """
    Notation AST statement that defines `lhs` as having the value `rhs`.
    """

    lhs: Variable
    rhs: NotationExpression

    @property
    def children(self):
        return [self.lhs, self.rhs]

Block dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that executes each of its arguments in turn.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Block(NotationTree, NotationStatement):
    """
    Notation AST statement that executes each of its arguments in turn.
    """

    bodies: tuple[NotationStatement, ...]

    @classmethod
    def from_children(cls, *bodies):
        return cls(bodies)

    @property
    def children(self):
        return list(self.bodies)

Cached dataclass

Bases: NotationTree, NotationExpression

Notation AST expression arg, equivalent to the quoted expression ref.

Often used after the compiler caches the computation ref into a variable arg, but we still wish to refer to the original expression to prove properties about it.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Cached(NotationTree, NotationExpression):
    """
    Notation AST expression `arg`, equivalent to the quoted expression `ref`.

    Often used after the compiler caches the computation `ref` into a variable
    `arg`, but we still wish to refer to the original expression to prove
    properties about it.
    """

    arg: NotationExpression
    ref: NotationExpression

    @property
    def result_type(self):
        return self.arg.result_type

    @property
    def children(self):
        return [self.arg, self.ref]

Call dataclass

Bases: NotationTree, NotationExpression, CallTerm

Notation AST expression for the result of calling the function op on args....

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Call(NotationTree, NotationExpression, CallTerm):
    """
    Notation AST expression for the result of calling the function `op` on
    `args...`.
    """

    op: NotationExpression
    args: tuple[NotationExpression, ...]

    @property
    def result_type(self) -> FType:
        arg_types = [a.result_type for a in self.args]
        op_type = self.op.result_type
        assert isinstance(op_type, CallableFType)
        return return_type(op_type, *arg_types)

    @classmethod
    def from_children(cls, op, *args):
        return cls(op, args)

    @property
    def children(self):
        return [self.op, *self.args]

Child dataclass

Bases: NotationTree, Cursor

A sublevel of a level expression.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Child(NotationTree, Cursor):
    """A sublevel of a level expression."""

    parent: Cursor
    attr: str = "lvl"

    @property
    def root(self):
        return self.parent.root

    @property
    def result_type(self):
        return self.parent.result_type.level_get_child_type(self.attr)

    @property
    def children(self):
        return [self.parent, Literal(self.attr)]

    @classmethod
    def from_children(cls, parent, attr):
        return cls(parent, attr.val)

Cursor dataclass

Bases: NotationExpression, ABC

A level expression that retains its owning tensor.

Source code in src/finch/finch_notation/nodes.py

class Cursor(NotationExpression, ABC):
    """A level expression that retains its owning tensor."""

    @property
    @abstractmethod
    def root(self) -> NotationExpression: ...

    @property
    @abstractmethod
    def result_type(self) -> LevelFType: ...

Declare dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that declares tns with an initial value init reduced with op in the current scope. A tensor with a dynamic fill takes init as its fill.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Declare(NotationTree, NotationStatement):
    """
    Notation AST statement that declares `tns` with an initial value `init` reduced
    with `op` in the current scope. A tensor with a dynamic fill takes `init` as
    its fill.
    """

    tns: NotationExpression
    init: NotationExpression
    op: NotationExpression
    shape: tuple[NotationExpression, ...]

    @property
    def children(self):
        return [self.tns, self.init, self.op, *self.shape]

    @classmethod
    def from_children(cls, tns, init, op, *shape):
        """
        Creates a Declare node from its children.
        """
        return cls(tns, init, op, shape)

from_children(tns, init, op, *shape) classmethod

Creates a Declare node from its children.

Source code in src/finch/finch_notation/nodes.py

@classmethod
def from_children(cls, tns, init, op, *shape):
    """
    Creates a Declare node from its children.
    """
    return cls(tns, init, op, shape)

Dimension dataclass

Bases: NotationTree, NotationExpression

Notation AST expression representing the dimension of tensor tns in rank r.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Dimension(NotationTree, NotationExpression):
    """
    Notation AST expression representing the dimension of tensor `tns` in
    rank `r`.
    """

    tns: NotationExpression
    r: Literal

    @property
    def result_type(self):
        assert isinstance(self.tns.result_type, TensorFType)
        return self.tns.result_type.shape_type[self.r.val]

    @classmethod
    def from_children(cls, tns, r):
        return cls(tns, r)

    @property
    def children(self):
        return [self.tns, self.r]

Fiber dataclass

Bases: NotationTree, NotationExpression

A positioned level view, retaining the owning tensor type at the root.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Fiber(NotationTree, NotationExpression):
    """A positioned level view, retaining the owning tensor type at the root."""

    lvl: Cursor
    pos: NotationExpression
    idxs: tuple[NotationExpression, ...] = ()

    @property
    def result_type(self):
        match self.lvl:
            case Root(tns):
                return tns.result_type
        root_type = self.lvl.root.result_type
        assert isinstance(root_type, TensorFType)
        return tensor.FiberTensorFType(self.lvl.result_type, root_type.device)

    @property
    def children(self):
        return [self.lvl, self.pos, *self.idxs]

    @classmethod
    def from_children(cls, lvl, pos, *idxs):
        return cls(lvl, pos, idxs)

FillValue dataclass

Bases: NotationTree, NotationExpression

Notation AST expression representing the fill value of tensor tns.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class FillValue(NotationTree, NotationExpression):
    """
    Notation AST expression representing the fill value of tensor `tns`.
    """

    tns: NotationExpression

    @property
    def result_type(self):
        assert isinstance(self.tns.result_type, TensorFType)
        return self.tns.result_type.element_type

    @classmethod
    def from_children(cls, tns):
        return cls(tns)

    @property
    def children(self):
        return [self.tns]

Freeze dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that freezes tns in the current scope after modifications with op.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Freeze(NotationTree, NotationStatement):
    """
    Notation AST statement that freezes `tns` in the current scope after
    modifications with `op`.
    """

    tns: NotationExpression
    op: NotationExpression

    @property
    def children(self):
        return [self.tns, self.op]

Full dataclass

Bases: NotationTree, NotationExpression

A read-only tensor filled with val, with scalar shape by default.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Full(NotationTree, NotationExpression):
    """A read-only tensor filled with `val`, with scalar shape by default."""

    val: NotationExpression
    shape: tuple[NotationExpression, ...] = ()

    @property
    def result_type(self):
        match self.val:
            case Literal(val):
                fill_type = StaticFillFType(val)
            case _:
                fill_type = DynamicFillFType(self.val.result_type)
        return tensor.patterns.FillTensorFType(
            fill_type,
            self.val.result_type,
            tuple(dim.result_type for dim in self.shape),
        )

    @classmethod
    def from_children(cls, val, *shape):
        return cls(val, tuple(shape))

    @property
    def children(self):
        return [self.val, *self.shape]

Function dataclass

Bases: NotationTree

A module-level function definition with arguments args....

Attributes:

Name	Type	Description
name	Variable	The function variable, annotated with its AssemblyKernelFType.
args	tuple[Variable, ...]	The arguments to the function.
body	NotationStatement	The body of the function. If it does not contain a return statement, the function returns the value of body.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Function(NotationTree):
    """
    A module-level function definition with arguments `args...`.

    Attributes:
        name: The function variable, annotated with its AssemblyKernelFType.
        args: The arguments to the function.
        body: The body of the function. If it does not contain a return statement,
            the function returns the value of `body`.
    """

    name: Variable
    args: tuple[Variable, ...]
    body: NotationStatement

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.name, *self.args, self.body]

    @classmethod
    def from_children(cls, name, *args_body):
        """Creates a term with the given head and arguments."""
        *args, body = args_body
        return cls(name, tuple(args), body)

children property

Returns the children of the node.

from_children(name, *args_body) classmethod

Creates a term with the given head and arguments.

Source code in src/finch/finch_notation/nodes.py

@classmethod
def from_children(cls, name, *args_body):
    """Creates a term with the given head and arguments."""
    *args, body = args_body
    return cls(name, tuple(args), body)

HollowFiber dataclass

Bases: Fiber

A candidate sparse fiber, recording whether a non-fill value was written.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class HollowFiber(Fiber):
    """A candidate sparse fiber, recording whether a non-fill value was written."""

    dirty: Variable = field(kw_only=True)

    @property
    def children(self):
        return [self.lvl, self.pos, self.dirty, *self.idxs]

    @classmethod
    def from_children(cls, lvl, pos, dirty, *idxs):
        return cls(lvl, pos, idxs, dirty=dirty)

If dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that only executes body if cond is true.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class If(NotationTree, NotationStatement):
    """
    Notation AST statement that only executes `body` if `cond` is true.
    """

    cond: NotationExpression
    body: NotationStatement

    @property
    def children(self):
        return [self.cond, self.body]

IfElse dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that executes then_body if cond is true, otherwise executes else_body.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class IfElse(NotationTree, NotationStatement):
    """
    Notation AST statement that executes `then_body` if `cond` is true, otherwise
    executes `else_body`.
    """

    cond: NotationExpression
    then_body: NotationStatement
    else_body: NotationStatement

    @property
    def children(self):
        return [self.cond, self.then_body, self.else_body]

Increment dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that updates the value lhs using rhs.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Increment(NotationTree, NotationStatement):
    """
    Notation AST statement that updates the value `lhs` using `rhs`.
    """

    lhs: Access
    rhs: NotationExpression

    @property
    def children(self):
        return [self.lhs, self.rhs]

Literal dataclass

Bases: NotationExpression, LiteralTerm

Notation AST expression for the literal value val.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Literal(NotationExpression, LiteralTerm):
    """
    Notation AST expression for the literal value `val`.
    """

    @property
    def result_type(self):
        return ftype(self.val)

    def __repr__(self) -> str:
        return literal_repr(type(self).__name__, {"val": self.val})

Loop dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that runs body for each value of idx in ext.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Loop(NotationTree, NotationStatement):
    """
    Notation AST statement that runs `body` for each value of `idx` in `ext`.
    """

    idx: Variable
    ext: NotationExpression
    body: NotationStatement

    @property
    def children(self):
        return [self.idx, self.ext, self.body]

Looplet dataclass

Bases: NotationExpression

A typed tensor expression being lowered by a looplet pass.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Looplet(NotationExpression):
    """A typed tensor expression being lowered by a looplet pass."""

    body: LoopletImpl
    type_: TensorFType

    @property
    def result_type(self) -> TensorFType:
        return self.type_

Module dataclass

Bases: NotationTree

Represents a group of functions. This is the toplevel translation unit for FinchNotation.

Attributes:

Name	Type	Description
funcs	tuple[Function, ...]	The functions defined in the module.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Module(NotationTree):
    """
    Represents a group of functions. This is the toplevel translation unit for
    FinchNotation.

    Attributes:
        funcs: The functions defined in the module.
    """

    funcs: tuple[Function, ...]

    @property
    def children(self):
        """Returns the children of the node."""
        return [*self.funcs]

    @classmethod
    def from_children(cls, *funcs):
        return cls(funcs)

children property

Returns the children of the node.

NotationCFGBuilder

Incrementally builds control-flow graph for Finch Notation IR.

Source code in src/finch/finch_notation/cfg_builder.py

class NotationCFGBuilder:
    """Incrementally builds control-flow graph for Finch Notation IR."""

    def __init__(self) -> None:
        self.cfg: ControlFlowGraph = ControlFlowGraph()
        self.current_block: BasicBlock = self.cfg.entry_block

    def emit(self, stmt) -> None:
        self.current_block.add_statement(stmt)

    def build(self, node: NotationNode) -> ControlFlowGraph:
        return self(node)

    def __call__(
        self,
        node: NotationNode,
        break_block: BasicBlock | None = None,
        return_block: BasicBlock | None = None,
    ) -> ControlFlowGraph:
        match node:
            case NumberedStatement(stmt, _):
                match stmt:
                    case Return(_):
                        self.emit(node)
                        assert return_block
                        self.current_block.add_successor(return_block)
                        unreachable_block = self.cfg.new_block()
                        self.current_block = unreachable_block
                    case _:
                        self.emit(node)
            case Return(_):
                # Allow building CFGs from unnumbered ASTs.
                self.emit(node)
                assert return_block
                self.current_block.add_successor(return_block)
                unreachable_block = self.cfg.new_block()
                self.current_block = unreachable_block
            case Block(bodies):
                for body in bodies:
                    self(body, break_block, return_block)
            case (
                Assign()
                | Increment()
                | Unpack()
                | Repack()
                | Declare()
                | Freeze()
                | Thaw()
            ):
                self.emit(node)
            case If(cond, body):
                # Treat `If` as `IfElse(cond, body, Block(()))`.
                self(IfElse(cond, body, Block(())), break_block, return_block)
            case IfElse(cond, then_body, else_body):
                before_block = self.current_block
                self.emit(cond)

                # create blocks for if, else, and after
                if_block = self.cfg.new_block()
                else_block = self.cfg.new_block()
                after_block = self.cfg.new_block()

                # connect before block to if and else blocks
                before_block.add_successor(if_block)
                before_block.add_successor(else_block)

                # fill in the if block
                self.current_block = if_block
                self(then_body, break_block, return_block)
                self.current_block.add_successor(after_block)

                # fill in the else block
                self.current_block = else_block
                self(else_body, break_block, return_block)
                self.current_block.add_successor(after_block)

                # continue building after the if-else
                self.current_block = after_block
            case Loop(_, ext, body):
                before_block = self.current_block
                # Record the loop extent expression (data dependency).
                self.emit(ext)

                # create blocks for the loop body and the code after the loop
                body_block = self.cfg.new_block()
                after_block = self.cfg.new_block()

                # connect before block to the loop body and the after block
                before_block.add_successor(body_block)
                before_block.add_successor(after_block)

                # fill in the loop body
                self.current_block = body_block
                self(body, after_block, return_block)

                # connect the end of loop body back to the beginning to form the loop
                self.current_block.add_successor(body_block)
                self.current_block.add_successor(after_block)
                self.current_block = after_block
            case Function(_, _, body):
                # Function argument definitions are inserted by `notation_desugar`.
                self(body, break_block, return_block)
            case Module(funcs):
                for func in funcs:
                    if not isinstance(func, Function):
                        raise NotImplementedError(
                            f"Unrecognized function type: {type(func)}"
                        )

                    if not isinstance(func.name, Variable):
                        raise NotImplementedError(
                            f"Unrecognized function name type: {type(func.name)}"
                        )

                    func_name = func.name.name

                    # set block names to the function name
                    self.cfg.block_name = func_name

                    # create entry/exit block for the function
                    func_entry_block = self.cfg.new_block()
                    func_exit_block = self.cfg.new_block()

                    # connect CFG entry block to the function's entry block
                    self.cfg.entry_block.add_successor(func_entry_block)

                    # dive into the body of the function
                    self.current_block = func_entry_block
                    self(func, break_block, func_exit_block)

                    # connect last block in the function to the
                    # exit block of the function
                    self.current_block.add_successor(func_exit_block)

                    # connect function exit block to the exit block of the CFG
                    func_exit_block.add_successor(self.cfg.exit_block)
            case other:
                raise NotImplementedError(other, "NotationCFGBuilder")

        return self.cfg

NotationExpression dataclass

Bases: NotationNode, ExpressionTerm

Notation AST expression base class.

A Notation expression is a program node which evaluates to a value.

Source code in src/finch/finch_notation/nodes.py

class NotationExpression(NotationNode, ExpressionTerm):
    """
    Notation AST expression base class.

    A Notation expression is a program node which evaluates to a value.
    """

    @property
    @abstractmethod
    def result_type(self) -> FType:
        """
        Get the type of the expression.
        """
        ...

result_type abstractmethod property

Get the type of the expression.

NotationInterpreter

Bases: UnvalidatedForm, NotationLoader

An interpreter for FinchNotation.

Source code in src/finch/finch_notation/interpreter.py

class NotationInterpreter(UnvalidatedForm, NotationLoader):
    """
    An interpreter for FinchNotation.
    """

    def __init__(
        self,
        bindings: ScopedDict | None = None,
        slots: ScopedDict | None = None,
        types: ScopedDict[FType] | None = None,
        loop_state: HaltState | None = None,
        function_state: HaltState | None = None,
    ):
        if bindings is None:
            bindings = ScopedDict()
        if slots is None:
            slots = ScopedDict()
        if types is None:
            types = ScopedDict()
        self.bindings = bindings
        self.slots = slots
        self.types = types
        self.loop_state = loop_state
        self.function_state = function_state

    def scope(
        self,
        bindings=None,
        slots=None,
        types=None,
        loop_state=None,
        function_state=None,
    ):
        """
        Create a new scope for the interpreter.
        This allows for nested scopes and variable shadowing.
        """
        if bindings is None:
            bindings = self.bindings.scope()
        if slots is None:
            slots = self.slots.scope()
        if types is None:
            types = self.types.scope()
        if loop_state is None:
            loop_state = self.loop_state
        if function_state is None:
            function_state = self.function_state
        return NotationInterpreter(
            bindings=bindings,
            slots=slots,
            types=types,
            loop_state=loop_state,
            function_state=function_state,
        )

    def lower(self, term: ntn.Module):
        return self._dispatch(term)

    @overload
    def __call__(self, prgm: ntn.Module) -> NotationInterpreterLibrary: ...

    @overload
    def __call__(self, prgm: ntn.NotationNode | asm.AssemblyNode) -> Any: ...

    def __call__(self, prgm: ntn.NotationNode | asm.AssemblyNode):
        """
        Run the program.
        """
        if isinstance(prgm, ntn.Module):
            return super().__call__(prgm)
        return self._dispatch(prgm)

    def _dispatch(self, prgm: ntn.NotationNode | asm.AssemblyNode):
        match prgm:
            case ntn.Literal(val) | asm.Literal(val):
                return val
            case ntn.Value(val, val_t):
                val_e = self(val)
                if ftype(val_e) is not val_t:
                    raise TypeError(
                        f"Value '{val_e}' is expected to be of ftype {val_t}, "
                        f"but is a type {ftype(val_e)}."
                    )
                return val_e
            case ntn.Variable(var_n, var_t):
                if var_n in self.types:
                    def_t = self.types[var_n]
                    if def_t != var_t:
                        raise TypeError(
                            f"Variable '{var_n}' is declared as type {def_t}, "
                            f"but used as type {var_t}."
                        )
                if var_n in self.bindings:
                    return self.bindings[var_n]
                raise KeyError(
                    f"Variable '{var_n}' is not defined in the current context."
                )
            case ntn.Call(f, args):
                f_e = self(f)
                args_e = [self(arg) for arg in args]
                return f_e(*args_e)
            case ntn.Root(tns):
                return self(tns).lvl
            case ntn.Child(parent, attr):
                return getattr(self(parent), attr)
            case ntn.Full(val, shape):
                return FullView(
                    self(val), tuple(self(dim) for dim in shape), prgm.result_type
                )
            case ntn.Unwrap(tns):
                return unwrap(self(tns))
            case ntn.Assign(var, val):
                val_e = self(val)
                if isinstance(var, ntn.Variable):
                    if var.type_ is not None:
                        assert fisinstance(val_e, var.type_)
                    var_n = var.name
                    self.bindings[var_n] = val_e
                    return None
                raise NotImplementedError(f"Unrecognized assignment target: {var}")
            case ntn.Slot(var_n, var_t):
                if var_n in self.types:
                    def_t = self.types[var_n]
                    if def_t != var_t:
                        raise TypeError(
                            f"Slot '{var_n}' is declared as type {def_t}, "
                            f"but used as type {var_t}."
                        )
                if var_n in self.slots:
                    return self.slots[var_n]
                raise KeyError(f"Slot '{var_n}' is not defined in the current context.")
            case ntn.Unpack(ntn.Slot(var_n, var_t), val):
                val_e = self(val)
                if not fisinstance(val_e, var_t):
                    raise TypeError(
                        f"Assigned value {val_e} is not of type {var_t} for "
                        f"variable '{var_n}'."
                    )
                assert var_n not in self.types, (
                    f"Variable '{var_n}' is already defined in the current"
                    f" context, cannot overwrite with slot."
                )
                self.types[var_n] = var_t
                self.slots[var_n] = val_e
                return None
            case ntn.Repack(ntn.Slot(var_n, var_t), val):
                if isinstance(val, ntn.Variable):
                    val_n = val.name
                    self.bindings[val_n] = self.slots[var_n]
                    del self.types[var_n]
                    del self.slots[var_n]
                    return None
                raise NotImplementedError(f"Unrecognized repack obj target: {val}")
            case ntn.Access(tns, mode, idxs):
                tns_e = self(tns)
                idxs_e = [self(idx) for idx in idxs]
                match mode:
                    case ntn.Read():
                        return access(tns_e, idxs_e)
                    case ntn.Update(op):
                        op_e = self(op)
                        return access(tns_e, idxs_e, op=op_e)
                    case _:
                        raise NotImplementedError(f"Unrecognized access mode: {mode}")

            case ntn.Dimension(tns, r):
                tns_e = self(tns)
                r_e = self(r)
                shape_ft = tns_e.ftype.shape_type[r_e]
                return shape_ft(tns_e.shape[r_e])
            case ntn.FillValue(tns):
                return self(tns).fill_value
            case ntn.Increment(tns, val):
                tns_e = self(tns)
                val_e = self(val)
                increment(tns_e, val_e)
                return None
            case ntn.Block(bodies):
                for body in bodies:
                    self(body)
                return None
            case ntn.Loop(idx, ext, body):
                ext_e = self(ext)
                # assert isinstance(ext_e, Extent)
                ext_e.loop_interpreter(self, idx, body)
                return None
            case ntn.Declare(tns, init, op, shape):
                assert isinstance(tns, ntn.Slot)
                tns_e = self(tns)
                init_e = self(init)
                op_e = self(op)
                shape_e = [self(s) for s in shape]
                self.slots[tns.name] = declare(tns_e, init_e, op_e, shape_e)
                return None
            case ntn.Freeze(tns, op):
                assert isinstance(tns, ntn.Slot)
                tns_e = self(tns)
                op_e = self(op)
                self.slots[tns.name] = freeze(tns_e, op_e)
                return None
            case ntn.Thaw(tns, op):
                assert isinstance(tns, ntn.Slot)
                tns_e = self(tns)
                op_e = self(op)
                self.slots[tns.name] = thaw(tns_e, op_e)
                return None
            case ntn.If(cond, body):
                if self(cond):
                    ctx_2 = self.scope()
                    ctx_2(body)
                return None
            case ntn.IfElse(cond, body, else_body):
                if not self(cond):
                    body = else_body
                ctx_2 = self.scope()
                ctx_2(body)
                return None
            case ntn.Function(
                ntn.Variable(
                    func_n, asm.AssemblyKernelFType(result_type=ret_t) as func_type
                ),
                args,
                body,
            ):

                def my_func(*args_e):
                    ctx_2 = self.scope(function_state=HaltState())
                    if len(args_e) != len(args):
                        raise ValueError(
                            f"Function '{func_n}' expects {len(args)} arguments, "
                            f"but got {len(args_e)}."
                        )
                    for arg, arg_e in zip(args, args_e, strict=False):
                        match arg:
                            case ntn.Variable(arg_n, _):
                                ctx_2.bindings[arg_n] = arg_e
                            case _:
                                raise NotImplementedError(
                                    f"Unrecognized argument type: {arg}"
                                )
                    ctx_2(body)
                    if ctx_2.function_state.has_returned:
                        ret_e = ctx_2.function_state.return_value
                        if not fisinstance(ret_e, ret_t):
                            raise TypeError(
                                f"Return value {ret_e} is not of type {ret_t} "
                                f"for function '{func_n}'. have type {ftype(ret_e)}"
                            )
                        return ret_e
                    raise ValueError(
                        f"Function '{func_n}' did not return a value, "
                        f"but expected type {ret_t}."
                    )

                self.bindings[func_n] = NotationInterpreterKernel(my_func, func_type)
                self.types[func_n] = func_type
                return None
            case ntn.Return(value):
                assert self.function_state is not None
                self.function_state.has_returned = True
                self.function_state.return_value = self(value)
                return None
            case ntn.Module(funcs):
                ctx_2 = self.scope()
                for func in funcs:
                    ctx_2(func)
                kernels = {}
                for func in funcs:
                    match func:
                        case ntn.Function(
                            ntn.Variable(
                                func_n,
                                ret_t,
                            ),
                            args,
                            _,
                        ):
                            kernels[func_n] = ctx_2.bindings[func_n]
                        case _:
                            raise NotImplementedError(
                                f"Unrecognized function definition: {func}"
                            )
                return NotationInterpreterLibrary(ctx_2, kernels)
            case _:
                raise NotImplementedError(
                    f"Unrecognized notation node type: {type(prgm)}"
                )

__call__(prgm)

__call__(prgm: ntn.Module) -> NotationInterpreterLibrary

__call__(prgm: ntn.NotationNode | asm.AssemblyNode) -> Any

Run the program.

Source code in src/finch/finch_notation/interpreter.py

def __call__(self, prgm: ntn.NotationNode | asm.AssemblyNode):
    """
    Run the program.
    """
    if isinstance(prgm, ntn.Module):
        return super().__call__(prgm)
    return self._dispatch(prgm)

scope(bindings=None, slots=None, types=None, loop_state=None, function_state=None)

Create a new scope for the interpreter. This allows for nested scopes and variable shadowing.

Source code in src/finch/finch_notation/interpreter.py

def scope(
    self,
    bindings=None,
    slots=None,
    types=None,
    loop_state=None,
    function_state=None,
):
    """
    Create a new scope for the interpreter.
    This allows for nested scopes and variable shadowing.
    """
    if bindings is None:
        bindings = self.bindings.scope()
    if slots is None:
        slots = self.slots.scope()
    if types is None:
        types = self.types.scope()
    if loop_state is None:
        loop_state = self.loop_state
    if function_state is None:
        function_state = self.function_state
    return NotationInterpreter(
        bindings=bindings,
        slots=slots,
        types=types,
        loop_state=loop_state,
        function_state=function_state,
    )

NotationLoader

Bases: Stage

Source code in src/finch/finch_notation/stages.py

class NotationLoader(Stage):
    @abstractmethod
    def lower(self, term: ntn.Module) -> asm.AssemblyLibrary:
        """
        Load the given notation program into a runnable module.
        """

lower(term) abstractmethod

Load the given notation program into a runnable module.

Source code in src/finch/finch_notation/stages.py

@abstractmethod
def lower(self, term: ntn.Module) -> asm.AssemblyLibrary:
    """
    Load the given notation program into a runnable module.
    """

NotationNode dataclass

Bases: Term, ABC

NotationNode

Base class for all Finch Notation nodes

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class NotationNode(Term, ABC):
    """
    NotationNode

    Base class for all Finch Notation nodes
    """

    @classmethod
    def head(cls):
        """Returns the head of the node."""
        return cls

    @classmethod
    def make_term(cls, head, *children):
        return head.from_children(*children)

    @classmethod
    def from_children(cls, *children):
        return cls(*children)

    def __str__(self):
        """Returns a string representation of the node."""
        ctx = NotationPrinterContext()
        res = ctx(self)
        return res if res is not None else ctx.emit()

__str__()

Returns a string representation of the node.

Source code in src/finch/finch_notation/nodes.py

def __str__(self):
    """Returns a string representation of the node."""
    ctx = NotationPrinterContext()
    res = ctx(self)
    return res if res is not None else ctx.emit()

head() classmethod

Returns the head of the node.

Source code in src/finch/finch_notation/nodes.py

@classmethod
def head(cls):
    """Returns the head of the node."""
    return cls

NotationStatement dataclass

Bases: NotationNode

Notation AST statement base class.

A Notation statement is a program node nested inside a function which does not produce a value, but may modify the state of the machine.

Source code in src/finch/finch_notation/nodes.py

class NotationStatement(NotationNode):
    """
    Notation AST statement base class.

    A Notation statement is a program node nested inside a function which does
    not produce a value, but may modify the state of the machine.
    """

NotationTransform

Bases: Stage

Source code in src/finch/finch_notation/stages.py

class NotationTransform(Stage):
    @abstractmethod
    def lower(self, term: ntn.Module) -> ntn.Module:
        """
        Transform the given assembly term into another assembly term.
        """

lower(term) abstractmethod

Transform the given assembly term into another assembly term.

Source code in src/finch/finch_notation/stages.py

@abstractmethod
def lower(self, term: ntn.Module) -> ntn.Module:
    """
    Transform the given assembly term into another assembly term.
    """

NumberedStatement dataclass

Bases: NotationStatement

Wrapper for NotationStatement that assigns a unique id to each statement.

Source code in src/finch/finch_notation/cfg_builder.py

@dataclass(eq=True, frozen=True)
class NumberedStatement(NotationStatement):
    """Wrapper for NotationStatement that assigns a unique id to each statement."""

    stmt: NotationStatement
    sid: int

    def __str__(self) -> str:
        return f"[{self.sid}] {str(self.stmt)}"

    @property
    def children(self):
        return (self.stmt, self.sid)

Read dataclass

Bases: AccessMode

Notation AST node representing a read-only access mode for a tensor. This mode allows reading the value of a tensor without modifying it.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Read(AccessMode):
    """
    Notation AST node representing a read-only access mode for a tensor.
    This mode allows reading the value of a tensor without modifying it.
    """

    @property
    def children(self):
        return []

Repack dataclass

Bases: NotationTree, NotationStatement

Registers updates from a symbolic object val with the original object obj. The original object may now be accessed and modified.

Attributes:

Name	Type	Description
slot		The symbolic object to read from.
obj	NotationExpression	The original object to write to.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Repack(NotationTree, NotationStatement):
    """
    Registers updates from a symbolic object `val` with the original
    object `obj`. The original object may now be accessed and modified.

    Attributes:
        slot: The symbolic object to read from.
        obj: The original object to write to.
    """

    val: Slot
    obj: NotationExpression

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.val, self.obj]

children property

Returns the children of the node.

Return dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that returns the value of val from the current function.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Return(NotationTree, NotationStatement):
    """
    Notation AST statement that returns the value of `val` from the current
    function.
    """

    val: NotationExpression

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.val]

    @classmethod
    def from_children(cls, val):
        return cls(val)

children property

Returns the children of the node.

Root dataclass

Bases: NotationTree, Cursor

The first level of a tensor.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Root(NotationTree, Cursor):
    """The first level of a tensor."""

    tns: NotationExpression

    @property
    def root(self):
        return self.tns

    @property
    def result_type(self):
        return cast("FinchTensorFType", self.tns.result_type).get_child_type("lvl")

    @property
    def children(self):
        return [self.tns]

Slot dataclass

Bases: NotationExpression, NamedTerm

Represents a register to a symbolic object. Using a register in an expression creates a copy of the object.

Attributes:

Name	Type	Description
name	str	The name of the symbolic object to register.
type	FType	The type of the symbolic object.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Slot(NotationExpression, NamedTerm):
    """
    Represents a register to a symbolic object. Using a register in an
    expression creates a copy of the object.

    Attributes:
        name: The name of the symbolic object to register.
        type: The type of the symbolic object.
    """

    name: str
    type: FType

    @property
    def result_type(self):
        """Returns the type of the expression."""
        return self.type

    def __repr__(self) -> str:
        return literal_repr(type(self).__name__, {"name": self.name, "type": self.type})

    @property
    def symbol(self) -> str:
        return self.name

result_type property

Returns the type of the expression.

TensorView

Bases: Tensor[FT]

Source code in src/finch/finch_notation/interpreter.py

class TensorView(Tensor[FT]):
    def __init__(self, idxs: tuple[Any, ...], tns: Tensor[FT], op: Any | None = None):
        """
        Initialize the TensorView with the specified indices, tensor, and operation.

        :param idxs: Tuple of index types for the tensor view.
        :param tns: The underlying tensor node.
        :param op: The operation applied to the tensor view, if any.
        """
        self.idxs = idxs
        self.tns = tns
        self.op = op

    @property
    def ftype(self):
        """
        Get the ftype of the tensor view.
        This is the ftype of the tensor at the specified indices.
        """
        return TensorViewFType(
            tuple(ftype(i) for i in self.idxs), self.tns.ftype, self.op
        )

    @property
    def shape(self):
        """
        Get the shape of the tensor view.
        This is the shape of the tensor at the specified indices.
        """
        return self.tns.shape[len(self.idxs) : -1]

    @property
    def ndim(self):
        """
        Get the number of dimensions of the tensor view.
        """
        return len(self.shape)

    @property
    def element_type(self):
        """
        Get the element type of the tensor view.
        This is the type of the elements in the tensor at the specified indices.
        """
        return self.tns.ftype.element_type

    @property
    def fill_value(self):
        """
        Get the fill value of the tensor view.
        This is the value used to fill the tensor at the specified indices.
        """
        return self.tns.fill_value

    @property
    def shape_type(self):
        """
        Get the shape type of the tensor view.
        This is the shape type of the tensor at the specified indices.
        """
        return self.tns.ftype.shape_type[len(self.idxs) : -1]

    def access(self, idxs, op=None):
        """
        Unfurl the tensor view along a specific index.
        This creates a new tensor view with the specified index unfurled.
        """
        return TensorView(idxs=self.idxs + idxs, tns=self.tns, op=op)

    def item(self):
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to Python scalar.")
        return self.unwrap()

    def to_numpy(self):
        return np.asarray(self.tns.to_numpy())[self.idxs]

    def to_scipy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_scipy.")

    def unwrap(self):
        """
        Unwrap the tensor view to get a scalar.
        """
        val = self.tns[*self.idxs]  # ty: ignore[not-subscriptable]
        assert isinstance(val, Tensor) and val.ndim == 0
        return val.item()

    def increment(self, val):
        """
        Increment the value in the tensor view.
        This updates the tensor at the specified index with the operation and value.
        """
        lhs = self.tns[*self.idxs]  # ty: ignore[not-subscriptable]
        assert isinstance(lhs, Tensor) and lhs.ndim == 0
        lhs = lhs.item()
        self.tns[*self.idxs] = self.op(lhs, val)  # ty: ignore[invalid-assignment]
        return

element_type property

Get the element type of the tensor view. This is the type of the elements in the tensor at the specified indices.

fill_value property

Get the fill value of the tensor view. This is the value used to fill the tensor at the specified indices.

ftype property

Get the ftype of the tensor view. This is the ftype of the tensor at the specified indices.

ndim property

Get the number of dimensions of the tensor view.

shape property

Get the shape of the tensor view. This is the shape of the tensor at the specified indices.

shape_type property

Get the shape type of the tensor view. This is the shape type of the tensor at the specified indices.

__init__(idxs, tns, op=None)

Initialize the TensorView with the specified indices, tensor, and operation.

:param idxs: Tuple of index types for the tensor view. :param tns: The underlying tensor node. :param op: The operation applied to the tensor view, if any.

Source code in src/finch/finch_notation/interpreter.py

def __init__(self, idxs: tuple[Any, ...], tns: Tensor[FT], op: Any | None = None):
    """
    Initialize the TensorView with the specified indices, tensor, and operation.

    :param idxs: Tuple of index types for the tensor view.
    :param tns: The underlying tensor node.
    :param op: The operation applied to the tensor view, if any.
    """
    self.idxs = idxs
    self.tns = tns
    self.op = op

access(idxs, op=None)

Unfurl the tensor view along a specific index. This creates a new tensor view with the specified index unfurled.

Source code in src/finch/finch_notation/interpreter.py

def access(self, idxs, op=None):
    """
    Unfurl the tensor view along a specific index.
    This creates a new tensor view with the specified index unfurled.
    """
    return TensorView(idxs=self.idxs + idxs, tns=self.tns, op=op)

increment(val)

Increment the value in the tensor view. This updates the tensor at the specified index with the operation and value.

Source code in src/finch/finch_notation/interpreter.py

def increment(self, val):
    """
    Increment the value in the tensor view.
    This updates the tensor at the specified index with the operation and value.
    """
    lhs = self.tns[*self.idxs]  # ty: ignore[not-subscriptable]
    assert isinstance(lhs, Tensor) and lhs.ndim == 0
    lhs = lhs.item()
    self.tns[*self.idxs] = self.op(lhs, val)  # ty: ignore[invalid-assignment]
    return

unwrap()

Unwrap the tensor view to get a scalar.

Source code in src/finch/finch_notation/interpreter.py

def unwrap(self):
    """
    Unwrap the tensor view to get a scalar.
    """
    val = self.tns[*self.idxs]  # ty: ignore[not-subscriptable]
    assert isinstance(val, Tensor) and val.ndim == 0
    return val.item()

Thaw dataclass

Bases: NotationTree, NotationStatement

Notation AST statement that thaws tns in the current scope, moving the tensor from read-only mode to update-only mode with a reduction operator op.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Thaw(NotationTree, NotationStatement):
    """
    Notation AST statement that thaws `tns` in the current scope, moving the tensor
    from read-only mode to update-only mode with a reduction operator `op`.
    """

    tns: NotationExpression
    op: NotationExpression

    @property
    def children(self):
        return [self.tns, self.op]

Unpack dataclass

Bases: NotationTree, NotationStatement

Attempts to convert rhs into a symbolic, which can be registerd with lhs. The original object must not be accessed or modified until the corresponding Repack node is reached.

Attributes:

Name	Type	Description
lhs	Slot	The symbolic object to write to.
rhs	NotationExpression	The original object to read from.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Unpack(NotationTree, NotationStatement):
    """
    Attempts to convert `rhs` into a symbolic, which can be registerd with
    `lhs`. The original object must not be accessed or modified until the
    corresponding `Repack` node is reached.

    Attributes:
        lhs: The symbolic object to write to.
        rhs: The original object to read from.
    """

    lhs: Slot
    rhs: NotationExpression

    @property
    def children(self):
        """Returns the children of the node."""
        return [self.lhs, self.rhs]

children property

Returns the children of the node.

Unwrap dataclass

Bases: NotationTree, NotationExpression

Notation AST statement that unwraps the scalar value from a 0-dimensional tensor arg.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Unwrap(NotationTree, NotationExpression):
    """
    Notation AST statement that unwraps the scalar value from a 0-dimensional
    tensor `arg`.
    """

    arg: NotationExpression

    @property
    def children(self):
        return [self.arg]

    @property
    def result_type(self):
        """
        Returns the type of the unwrapped value.
        """
        assert isinstance(self.arg.result_type, TensorFType)
        return self.arg.result_type.element_type

result_type property

Returns the type of the unwrapped value.

Update dataclass

Bases: AccessMode, NotationTree

Notation AST node representing an update access mode for a tensor. This mode allows reading and modifying the value of a tensor. Increment operations are allowed in this mode, and will use the update operation op to increment ref with val as ref = op(ref, val). To overwrite the value of a tensor, use the ops algebra.overwrite or algebra.init_write. Attributes: op: The operation used to update the value of the tensor.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Update(AccessMode, NotationTree):
    """
    Notation AST node representing an update access mode for a tensor.  This
    mode allows reading and modifying the value of a tensor.  Increment
    operations are allowed in this mode, and will use the update operation `op`
    to increment `ref` with `val` as `ref = op(ref, val)`.  To overwrite the
    value of a tensor, use the ops `algebra.overwrite` or `algebra.init_write`.
    Attributes:
        op: The operation used to update the value of the tensor.
    """

    op: NotationExpression

    @property
    def children(self):
        return [self.op]

Value dataclass

Bases: NotationExpression

Notation AST expression for host code val expected to evaluate to a value of type type_.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Value(NotationExpression):
    """
    Notation AST expression for host code `val` expected to evaluate to a value of
    type `type_`.
    """

    ex: AssemblyExpression
    type_: FType

    @property
    def result_type(self):
        return self.type_

    def __repr__(self) -> str:
        return literal_repr(type(self).__name__, {"ex": self.ex, "type_": self.type_})

Variable dataclass

Bases: NotationExpression, NamedTerm

Notation AST expression for a variable named name.

Attributes:

Name	Type	Description
name	str	The name of the variable.
type_	FType	The type of the variable.

Source code in src/finch/finch_notation/nodes.py

@dataclass(eq=True, frozen=True)
class Variable(NotationExpression, NamedTerm):
    """
    Notation AST expression for a variable named `name`.

    Attributes:
        name: The name of the variable.
        type_: The type of the variable.
    """

    name: str
    type_: FType

    def __post_init__(self):
        assert isinstance(self.type_, FType)

    @property
    def result_type(self):
        return self.type_

    def __repr__(self) -> str:
        return literal_repr(
            type(self).__name__, {"name": self.name, "type_": self.type_}
        )

    @property
    def symbol(self) -> str:
        return self.name