bdusell / semiring-einsum

Generic PyTorch implementation of einsum that supports different semirings

Semiring Einsum

This is an extensible PyTorch implementation of einsum that supports multiple semirings.

Read the full documentation here.

semiring-einsum's Issues

More than 26 indices?

torch.einsum limits you to 26 indices but can semiring_einsum provide a non-string interface that allows more than 26?

log_viterbi_einsum_forward raises exception with no summed-out variables

eq = compile_equation('a,a->a')
x = torch.arange(5, dtype=float)
y = torch.arange(5, dtype=float)

log_viterbi_einsum_forward(eq, x, y, block_size=1)

Expected output:
(tensor([0., 2., 4., 6., 8.], dtype=torch.float64), tensor([], size=(5, 0)))

Actual output:

  File ".../torch_semiring_einsum/log_viterbi_forward.py", line 102, in max_argmax_block
    argmax = torch.stack(argmaxes, dim=-1)
RuntimeError: stack expects a non-empty TensorList

Question about how to choose block_size

Hi, is there any guidance on how to choose block_size?

Allow repeated and unmatched output indices?

Although the original einsum doesn't allow it, it would be convenient if:

• einsum('i->ii', torch.tensor([1., 2.])) returns torch.tensor([[1., 0.], [0., 2.]]) (that is, it's like torch.diag_embed)
• einsum('->i', torch.tensor(1.) returns torch.tensor([1.]) (that is, it's like torch.unsqueeze)

I can imagine that the second proposal is more controversial than the first one.

log_einsum returns nan instead of inf

log_einsum(s.compile_equation(',->'), torch.tensor(2e38), torch.tensor(2e38), block_size=1)

Expected output: torch.tensor(inf)

Actual output: torch.tensor(nan)

Default block_size

Or can the documentation at least recommend a good block_size?

Boolean semiring

How hard would it be to add the Boolean semiring?

In log forward, include tensors saved for backward in size calculation

When calculating the space required for the output tensors in the forward pass of log einsum, the tensors of max values and sumexpsub values are currently not included, leading to a very low estimate. Use output_dtypes to include them.

Doesn't work with PyTorch 2?

The pyproject.toml (if I understand correctly) requires PyTorch >= 1.1 and < 2.0. Is that necessary?

`log_viterbi_einsum_forward` won't work if there are no summed variables

It looks like max_argmax_block needs to be modified so that it works if dims is empty.

sync examples regarding specifying block_size

First of all, thanks very much for this great work! (I'm surprised that I'm only now finding this since it it's been on my wish-list for years! A year or two ago, I also ran into the block_size trade off with my own generalized semiring tensordot implementation, so I'm happy to see that in here, too.)

Related to issue #1 , the examples on the home page don't specify a block_size which leads to a runtime error:

C = torch_semiring_einsum.log_einsum(EQUATION, A, B)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-33-d81285427569> in <module>
----> 1 C = torch_semiring_einsum.log_einsum(EQUATION, A, B)

TypeError: result() missing 1 required keyword-only argument: 'block_size'

Specifying block_size fixes it:

C = torch_semiring_einsum.log_einsum(EQUATION, A, B, block_size=1)

One of the following should probably be done:

1. update the examples to include the block_size
2. add a default value for the block_size (this is probably preferred since new users wouldn't need to figure out the block_size thing in order to get started using the package).

Even-better, I wonder about the feasibility of implementing a more reasonable, adaptive default.
From your explanation on #1 , It sounds like you have a reasonable heuristic in mind for block_size as a function of input tensor shapes and available memory (psutil can get access to available memory: e.g. psutil.virtual_memory().available). Maybe that function could be constructed and cached at the same time that the equation is compiled and then block_size=torch_semiring_einsum.heuristic_block_size could indicate that it should be used to compute the block_size? Is it worth opening a separate issue and possibly a pull-request?

(feat): Support For Complex Numbers

Hello! This seems like it would be a huge boost for me, but I need support for complex numbers - any chance I could make a PR? Right now I'm getting:

a = tensor([[[9.9567e-01, 9.9567e-01, 9.9567e-01,  ..., 9.9567e-01,
          9.9567e-01, 9.9567e-01],
         [9.9567e-0...       [1.6774e-05, 1.6774e-05, 1.6774e-05,  ..., 1.6774e-05,
          1.6774e-05, 1.6774e-05]]], dtype=torch.float64)
b = tensor([[[ 0.0000+0.0000j, -0.0024-0.0698j, -0.0097-0.1392j,  ...,
          -0.0219+0.2079j, -0.0097+0.1392j, -0.0024...698j, -0.0097-0.1392j,  ...,
          -0.0219+0.2079j, -0.0097+0.1392j, -0.0024+0.0698j]]],
       requires_grad=True)

    def multiply_in_place(a, b):
>       a.mul_(b)
E       RuntimeError: result type ComplexDouble can't be cast to the desired output type Double

Folding trick

Per https://twitter.com/BrianDuSell/status/1333272047858937856 .

Broadcasting

torch.einsum now supports broadcasting: pytorch/pytorch#48420

Better error messages

A common error I get is in Equation.validate_sizes():

IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)

which is not as clear as the other error message that this function prints out.

out= option

Can an out= option be added to overwite an existing tensor instead of creating a new one?

[feature request] support multiple simultaneous right-hand-sides

I have a situation where many of my formulas come in sets like the following:

'bijklmn,bij,bkl,bm,bn->'
'bijklmn,bij,bkl,bm,bn->bij'
'bijklmn,bij,bkl,bm,bn->bkl'
'bijklmn,bij,bkl,bm,bn->bm'
'bijklmn,bij,bkl,bm,bn->bn'

It seems like there should be a substantial amount of work between these that could be shared, but I might be misunderstanding the way that you are doing the reducing.

If there is a possibility to support more efficient bulk work (or maybe even if there isn't), I'd suggest allowing the above set of formulas to be replaced with the following (which should indicate that *einsum should return a tuple of tensors with the respective results):

'bijklmn,bij,bkl,bm,bn->,bij,bkl,bm,bn'

@bdusell would you please comment on the following when you get a chance:

1. Are you open to the proposed notation?
2. Are you open to pull requests?
3. Do you see an opportunity for reused computation in the above scenario? (after a quick look, it doesn't look as easy as I had initially thought)

Thanks!

Zero-dimensional tensors

Desired behavior:

>>> torch.einsum('->', torch.tensor(1.))
tensor(1.)

Actual behavior:

>>> eq = compile_equation('->')
>>> einsum(eq, torch.tensor(1.), block_size=1)
...
  File ".../torch_semiring_einsum/equation.py", line 159, in lookup
    return arg[index].permute(self.permutation)
IndexError: too many indices for tensor of dimension 0