Einsum in python

[willow-ahrens]

In addition or instead of supporting np.einsum, it was agreed that finch should have a more flexible einsum interface. It would look like this:

allocating scaled matrix add

einsum("d[i, j] += a[i, j]*b[] + c[i, j]", a=a, b=b, c=c)

in-place scaled matrix add

einsum("d[i, j] += a[i, j]*b[] + c[i, j]", a=a, b=b, c=c)

max-plus matmul

einsum("c[i, j] max= a[i, j] + b[i, j]", a=a, b=b)

etc. If you want to use your own function you could pass it in

einsum("c[i, j] f= a[i, j] + b[i, j]", a=a, b=b, f=f)

we would target https://github.com/willow-ahrens/Finch.jl/issues/428 and probably utilize https://github.com/lark-parser/lark or https://ply.readthedocs.io/en/latest/ or a custom recursive descent parser to parse the string.

@mtsokol @kylebd99

[hameerabbasi]

While einsum is great and we can (and should) support it, I'd suggest first supporting the functional forms of these operations. For examples given in the comment these would be:

d += a * b + c  # First two are identical
finch.max(c, a + b, out=c)
f(c, a + b, out=c)

The main reason is that the array API standard doesn't contain einsum, and we want users of the standard to be able to get the full performance benefits of fused operations. Quoting design principle 5:

an interface should have clearly defined and delineated functionality which, ideally, has no overlap with the functionality of other interfaces in the specification. Providing multiple interfaces which can all perform the same operation creates unnecessary confusion regarding interface applicability (i.e., which interface is best at which time) and decreases readability of both library and user code. Where overlap is possible, the specification must be parsimonious in the number of interfaces, ensuring that each interface provides a unique and compelling abstraction. Examples of related interfaces which provide distinct levels of abstraction (and generality) include ... einsum

[hameerabbasi]

For a discussion on existing ways to handle laziness in Python, see: https://github.com/pydata/sparse/discussions/618

[willow-ahrens]

That's okay, we'll work on the python side of this later! We wanted to agree on something that seemed doable because there was interest from the julia side in einsum and we were designing things to ensure that python wouldn't be locked out of it.

[willow-ahrens]

Kyle or I can tackle this after my next paper submission

[willow-ahrens]

@hameerabbasi , I'm curious how we can represent the following in the tensor API, it's important in graph kernels:

C[i, j] max= A[i, k] + B[k, j]

[hameerabbasi]

@hameerabbasi , I'm curious how we can represent the following in the tensor API, it's important in graph kernels:

C[i, j] max= A[i, k] + B[k, j]

Something like C = finch.max(A[:, None, :], finch.transpose(B, (1, 0))[None, :, :], axis=1)

An assumption made here is that A, B, C are all 2D and stored in the order i, j, k with k being the "contiguous" variable where present. Also, C doesn't exist beforehand. If that's not true, one would add/modify transpositions, or add another max with C like finch.max(C, ..., out=C).

[willow-ahrens]

okay. I still think the einsum is much clearer here, but I do see how this could eventually work through the fused interface, we would need indexing to be lazy too, which might be farther off.