nodes

Access dataclass

Bases: EinsumExpression, EinsumTree

Access

Tensor access like a[i, j].

Attributes:

Name	Type	Description
tensor		The tensor to access.
idxs	tuple[Index, ...]	The indices at which to access the tensor.

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

@dataclass(eq=True, frozen=True)
class Access(EinsumExpression, EinsumTree):
    """
    Access

    Tensor access like a[i, j].

    Attributes:
        tensor: The tensor to access.
        idxs: The indices at which to access the tensor.
    """

    tns: Alias
    idxs: tuple[Index, ...]  # (Field('i'), Field('j'))
    # Children: None (leaf)

    @classmethod
    def from_children(cls, *children: Term) -> Self:
        # First child is tns, rest are indices
        if len(children) < 1:
            raise ValueError("Access expects at least 1 child")
        tns = cast(Alias, children[0])
        idxs = cast(tuple[Index, ...], children[1:])
        return cls(tns, tuple(idxs))

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

    def get_idxs(self) -> set["Index"]:
        idxs = set()
        for idx in self.idxs:
            idxs.update(idx.get_idxs())
        return idxs

Alias dataclass

Bases: EinsumExpression

Represents a AST expression for an index named name.

Attributes:

Name	Type	Description
name	str	The name of the index.

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

@dataclass(eq=True, frozen=True)
class Alias(EinsumExpression):
    """
    Represents a  AST expression for an index named `name`.

    Attributes:
        name: The name of the index.
    """

    name: str

    def get_idxs(self) -> set["Index"]:
        return set()

Call dataclass

Bases: EinsumExpression, EinsumTree

Call

Represents an operation like + or * on pointwise expressions for multiple operands. If operation is not commutative, pointwise node must be binary, with 2 args at most.

Attributes:

Name	Type	Description
op	Literal	The function to apply e.g., ffunc.add, ffunc.mul, ffunc.subtract, ffunc.div, etc... Must be a callable.
args	tuple[EinsumExpression, ...]	The arguments to the operation.

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

@dataclass(eq=True, frozen=True)
class Call(EinsumExpression, EinsumTree):
    """
    Call

    Represents an operation like + or * on pointwise expressions for multiple operands.
    If operation is not commutative, pointwise node must be binary, with 2 args at most.

    Attributes:
        op: The function to apply e.g.,
            ffunc.add, ffunc.mul, ffunc.subtract, ffunc.div, etc...
            Must be a callable.
        args: The arguments to the operation.
    """

    op: Literal  # the function to apply e.g., ffunc.add
    args: tuple[EinsumExpression, ...]  # Subtrees
    # input_fields: tuple[tuple[Field, ...], ...]
    # Children: The args

    @classmethod
    def from_children(cls, *children: Term) -> Self:
        # First child is op, rest are args
        if len(children) < 2:
            raise ValueError("Call expects at least 2 children (op + 1 arg)")
        op = cast(Literal, children[0])
        args = cast(tuple[EinsumExpression, ...], children[1:])
        return cls(op, tuple(args))

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

    def get_idxs(self) -> set["Index"]:
        idxs = set()
        for arg in self.args:
            idxs.update(arg.get_idxs())
        return idxs

Einsum dataclass

Bases: EinsumTree, EinsumStatement

Einsum

A einsum operation that maps pointwise expressions and aggregates them.

Attributes:

Name	Type	Description
op	Literal	The function to apply to the pointwise expressions (e.g. +=, f=, max=, etc...).
idxs	tuple[Index, ...]	The indices that are used in the output (i.e. i, j).
arg	EinsumExpression	The pointwise expression that is mapped and aggregated.

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

@dataclass(eq=True, frozen=True)
class Einsum(EinsumTree, EinsumStatement):
    """
    Einsum

    A einsum operation that maps pointwise expressions and aggregates them.

    Attributes:
        op: The function to apply to the pointwise expressions
                    (e.g. +=, f=, max=, etc...).

        idxs: The indices that are used in the output
                    (i.e. i, j).

        arg: The pointwise expression that
                    is mapped and aggregated.
    """

    op: Literal
    tns: Alias
    idxs: tuple[Index, ...]
    arg: EinsumExpression

    @classmethod
    def from_children(cls, *children: Term) -> Self:
        # Expecting exactly 4 children
        if len(children) != 4:
            raise ValueError(f"Einsum expects 4 children, got {len(children)}")
        op = cast(Literal, children[0])
        tns = cast(Alias, children[1])
        idxs = cast(tuple[Index, ...], children[2])
        arg = cast(EinsumExpression, children[3])
        return cls(op, tns, idxs, arg)

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

EinsumExpression

Bases: EinsumNode, ABC

Einsum AST expression base class.

Represents a pointwise expression in the Einsum IR, which evaluates to a scalar.

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

class EinsumExpression(EinsumNode, ABC):
    """
    Einsum AST expression base class.

    Represents a pointwise expression in the Einsum IR, which evaluates to a scalar.
    """

    @abstractmethod
    def get_idxs(self) -> set["Index"]:
        pass

EinsumNode

Bases: Term

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

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

    @classmethod
    def make_term(cls, head, *children: Term) -> Self:
        return head.from_children(*children)

    @classmethod
    def from_children(cls, *children: Term) -> Self:
        return cls(*children)

    def __str__(self):
        """Returns a string representation of the node."""
        ctx = EinsumPrinterContext()
        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_einsum/nodes.py

def __str__(self):
    """Returns a string representation of the node."""
    ctx = EinsumPrinterContext()
    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_einsum/nodes.py

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

EinsumStatement

Bases: EinsumNode

Einsum AST statement base class.

Represents a statement in the Einsum IR. An Einsum statement may modify the state of the machine by assigning tensor values to variables. Einsum statements evaluate to a tuple of tensor values.

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

class EinsumStatement(EinsumNode):
    """
    Einsum AST statement base class.

    Represents a statement in the Einsum IR. An Einsum statement may modify
    the state of the machine by assigning tensor values to variables. Einsum
    statements evaluate to a tuple of tensor values.
    """

Index dataclass

Bases: EinsumExpression

Represents a AST expression for an index named name.

Attributes:

Name	Type	Description
name	str	The name of the index.

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

@dataclass(eq=True, frozen=True)
class Index(EinsumExpression):
    """
    Represents a  AST expression for an index named `name`.

    Attributes:
        name: The name of the index.
    """

    name: str

    def get_idxs(self) -> set["Index"]:
        return {self}

Literal dataclass

Bases: EinsumExpression

Literal

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

@dataclass(eq=True, frozen=True)
class Literal(EinsumExpression):
    """
    Literal
    """

    val: Any

    def __hash__(self):
        return hash(self.val)

    def __eq__(self, other):
        return isinstance(other, Literal) and self.val == other.val

    def get_idxs(self) -> set["Index"]:
        return set()

Plan dataclass

Bases: EinsumTree, EinsumStatement

Plan

A plan that contains einsum operations. Basically a list of einsums and some return values.

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

@dataclass(eq=True, frozen=True)
class Plan(EinsumTree, EinsumStatement):
    """
    Plan

    A plan that contains einsum operations.
    Basically a list of einsums and some return values.
    """

    bodies: tuple[EinsumStatement, ...] = ()

    @classmethod
    def from_children(cls, *children: Term) -> Self:
        # The last child is the returnValues tuple, all others are bodies
        if len(children) < 1:
            raise ValueError("Plan expects at least 1 child")
        bodies = children

        return cls(
            tuple(cast(EinsumStatement, b) for b in bodies),
        )

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

Produces dataclass

Bases: EinsumTree, EinsumStatement

Represents a logical AST statement that returns args... from the current plan. Halts execution of the program.

Attributes:

Name	Type	Description
args	tuple[Alias, ...]	The arguments to return.

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

@dataclass(eq=True, frozen=True)
class Produces(EinsumTree, EinsumStatement):
    """
    Represents a logical AST statement that returns `args...` from the current plan.
    Halts execution of the program.

    Attributes:
        args: The arguments to return.
    """

    args: tuple[Alias, ...]

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

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

children property

Returns the children of the node.