interpreter

FullView

Bases: Tensor

An interpreter view of a read-only tensor filled with one value.

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

class FullView(Tensor):
    """An interpreter view of a read-only tensor filled with one value."""

    def __init__(self, val, shape, type_):
        self.val = val
        self._shape = shape
        self._type = type_

    @property
    def ftype(self):
        return self._type

    @property
    def shape(self):
        return self._shape

    @property
    def fill_value(self):
        return self.val

    def access(self, idxs, op=None):
        if op is not None:
            raise TypeError("Full views are read-only")
        return FullView(
            self.val,
            self.shape[len(idxs) :],
            replace(self.ftype, _shape_type=self.ftype.shape_type[len(idxs) :]),
        )

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

    def unwrap(self):
        return self.item()

    def to_numpy(self):
        return np.full(self.shape, self.val)

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

HaltState dataclass

A class to represent the halt state of a notation program. These programs can't break, but calling return sets a special return value.

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

@dataclass(eq=True)
class HaltState:
    """
    A class to represent the halt state of a notation program.
    These programs can't break, but calling return sets a special return value.
    """

    has_returned: bool = False
    return_value: Any = None

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,
    )

NotationInterpreterKernel

Bases: AssemblyKernel

A kernel for interpreting FinchNotation code. This is a simple interpreter that executes the assembly code.

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

class NotationInterpreterKernel(asm.AssemblyKernel):
    """
    A kernel for interpreting FinchNotation code.
    This is a simple interpreter that executes the assembly code.
    """

    def __init__(self, func, type_):
        super().__init__(type_)
        self.func = func

    def __call__(self, *args):
        return self.func(*args)

NotationInterpreterLibrary

Bases: AssemblyLibrary

A class to represent an interpreted module of FinchNotation.

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

class NotationInterpreterLibrary(asm.AssemblyLibrary):
    """
    A class to represent an interpreted module of FinchNotation.
    """

    def __init__(self, ctx, kernels):
        self.ctx = ctx
        self.kernels = kernels

    def __getattr__(self, name):
        # Allow attribute access to kernels by name
        if name in self.kernels:
            return self.kernels[name]
        raise AttributeError(
            f"{type(self).__name__!r} object has no attribute {name!r}"
        )

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()

TensorViewFType

Bases: TensorFType[FT]

A ftype for tensor views. This is used to represent the ftype of a tensor at specific indices. It is a subclass of ntn.FType to allow for custom formatting.

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

class TensorViewFType(TensorFType[FT]):
    """
    A ftype for tensor views.
    This is used to represent the ftype of a tensor at specific indices.
    It is a subclass of ntn.FType to allow for custom formatting.
    """

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

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

    def __eq__(self, other):
        return (
            isinstance(other, TensorViewFType)
            and self.idxs == other.idxs
            and self.tns == other.tns
            and self.op == other.op
        )

    def __hash__(self):
        return hash((self.idxs, self.tns, self.op))

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

    @property
    def fill_type(self):
        """
        Get the fill of the tensor view's format.
        """
        return self.tns.fill_type

    @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.shape_type[len(self.idxs) : -1]

    def __call__(
        self, idxs: tuple[Any, ...], tns: Tensor[FT], op: Any | None = None
    ) -> TensorView:
        return TensorView(idxs=idxs, tns=tns, op=op)

    def construct(self, shape: tuple[Any, ...]) -> TensorView:
        raise RuntimeError("`TensorView`s cannot be constructed.")

    def from_numpy(self, arr: np.ndarray) -> TensorView:
        raise RuntimeError("`TensorView`s cannot be converted from NumPy arrays.")

fill_type property

Get the fill of the tensor view's format.

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 TensorViewFType with the specified indices, tensor, and operation.

:param idxs: Tuple of index types for the tensor view. :param tns: The underlying tensor ftype. :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: TensorFType[FT], op: Any = None):
    """
    Initialize the TensorViewFType with the specified indices, tensor, and
    operation.

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

access(tns, idxs, op=None)

Unfurl a tensor along an index. This is used to create a tensor view for a specific slice of the tensor.

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

def access(tns, idxs, op=None):
    """
    Unfurl a tensor along an index.
    This is used to create a tensor view for a specific slice of the tensor.
    """
    if hasattr(tns, "access"):
        return tns.access(idxs, op)
    return TensorView(idxs=idxs, tns=tns, op=op)

declare(tns, init, op, shape)

Declare a tensor.

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

def declare(tns, init, op, shape):
    """
    Declare a tensor.
    """
    if hasattr(tns, "declare"):
        return tns.declare(init, op, shape)
    match tns:
        case np.ndarray():
            return np_declare(tns, init, op, shape)
        case _:
            raise AttributeError(f"{type(tns).__name__} has no declare")

freeze(tns, op)

Freeze a tensor.

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

def freeze(tns, op):
    """
    Freeze a tensor.
    """
    if hasattr(tns, "freeze"):
        return tns.freeze(op)
    return tns

increment(tns, val)

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

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

def increment(tns, val):
    """
    Increment a tensor view with an operation and value.
    This updates the tensor at the specified index with the operation and value.
    """
    if hasattr(tns, "increment"):
        return tns.increment(val)
    raise AttributeError(f"{type(tns).__name__} has no increment")

thaw(tns, op)

Thaw a tensor.

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

def thaw(tns, op):
    """
    Thaw a tensor.
    """
    if hasattr(tns, "thaw"):
        return tns.thaw(op)
    return tns

unwrap(tns)

Unwrap a tensor view to get the underlying tensor. This is used to get the original tensor from a tensor view.

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

def unwrap(tns):
    """
    Unwrap a tensor view to get the underlying tensor.
    This is used to get the original tensor from a tensor view.
    """
    if hasattr(tns, "unwrap"):
        return tns.unwrap()
    raise AttributeError(f"{type(tns).__name__} has no unwrap")