Skip to content

finch_einsum

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 
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
@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 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
@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 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
@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 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
@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]

EinsumEvaluator

Bases: Stage

Source code in src/finch/finch_einsum/stages.py 
10
11
12
13
14
15
16
17
18
19
class EinsumEvaluator(Stage):
    @abstractmethod
    def lower(
        self,
        term: ein.EinsumNode,
        bindings: dict[ein.Alias, Any] | None = None,
    ):  # TODO eventually Any->Tensor
        """
        Evaluate the given logic.
        """

lower(term, bindings=None) abstractmethod

Evaluate the given logic.

Source code in src/finch/finch_einsum/stages.py 
11
12
13
14
15
16
17
18
19
@abstractmethod
def lower(
    self,
    term: ein.EinsumNode,
    bindings: dict[ein.Alias, Any] | None = None,
):  # TODO eventually Any->Tensor
    """
    Evaluate the given logic.
    """

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 
38
39
40
41
42
43
44
45
46
47
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

EinsumLoader

Bases: Stage

Source code in src/finch/finch_einsum/stages.py 
22
23
24
25
26
27
28
29
class EinsumLoader(Stage):
    @abstractmethod
    def lower(self, term: ein.EinsumStatement, bindings: dict[ein.Alias, TensorFType]):
        """
        Generate Finch Library from the given logic and input types, with a
        single method called `main` which implements the logic. Also return a
        dictionary including additional tables needed to run the kernel.
        """

lower(term, bindings) abstractmethod

Generate Finch Library from the given logic and input types, with a single method called main which implements the logic. Also return a dictionary including additional tables needed to run the kernel.

Source code in src/finch/finch_einsum/stages.py 
23
24
25
26
27
28
29
@abstractmethod
def lower(self, term: ein.EinsumStatement, bindings: dict[ein.Alias, TensorFType]):
    """
    Generate Finch Library from the given logic and input types, with a
    single method called `main` which implements the logic. Also return a
    dictionary including additional tables needed to run the kernel.
    """

EinsumNode

Bases: Term

Source code in src/finch/finch_einsum/nodes.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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 
24
25
26
27
28
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 
11
12
13
14
@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 
50
51
52
53
54
55
56
57
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 
78
79
80
81
82
83
84
85
86
87
88
89
90
@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 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
@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 
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
@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 
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
@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.