Rename symbols to support torch.einsum

[fritzo]

Description

This renames symbols to support PyTorch 0.4.1 which supports only symbols a-z. This is useful when the overall computation has more than 26 dimensions but any single contraction requires only 26 or fewer dimensions.

Questions

• Do other backends work? I have only tested PyTorch. If other backends don't work, I'll simply limit the test to the PyTorch backend.

Status

• Ready to merge as soon as tests pass

Tested

• Added a test with non-lowercase symbols
• Tested that this works in an application using Pyro Use opt_einsum for message passing in SVI pyro-ppl/pyro#1313

[dgasmith]

Overall LGTM. It looks like CI missed a trigger (very odd), can you make another commit to trigger CI when you get the chance.

[dgasmith]

Extreme, but can we throw if len(symbols) > 26.

[fritzo]

The advantage of not throwing is that this will continue to work when PyTorch fixes their einsum to allow more symbols. I'd rather let their current implementation throw.

[dgasmith]

Makes sense. I am not intimately familiar with the new torch symbols set, but would opt_einsum.parser.convert_to_valid_einsum_chars work here?

[dgasmith]

@jcmgray We do this for einsum as well I believe. Is it time to have a single expression which converts from global to local symbols?

[dgasmith]

@jcmgray Can you look over this as well?

[fritzo]

(Note the only reason I've removed this line is so I can reuse the einsum() wrapper in this file, rather than calling torch.einsum in two different places)

[jcmgray]

Do other backends work? I have only tested PyTorch. If other backends don't work, I'll simply limit the test to the PyTorch backend.

einsum in numpy, cupy, dask and tensorflow (from only recently) all support upper case letters so tests should be fine. The ideal situation would be for pytorch to add support for upper case, but the implementation is c++ so its not super obvious to me how easy that would be (@t-vi?). Probably easier to add it here anyway!

---

It might be cleaner, and more useful for potential future backends, to modify the existing machinery:

• opt_einsum.parser.convert_to_valid_einsum_chars
• opt_einsum.parser.has_valid_einsum_chars_only
• opt_einsum.parser.is_valid_einsum_char

probably with a allow_uppercase keyword, which could just switch is_valid_einsum_char between checking against:

import string
string.ascii_lowercase  # or
string.ascii_letters

Also, for what its worth, your strategy of just replacing all characters might be faster than replacing just the invalid ones, in which case do update convert_to_valid_einsum_chars with your snippet.

[t-vi]

PyTorch's einsum makes some use of the fact that valid letters are consecutive and from a small fixed set (based on review comments when I previously dynamically allocated indices). It's not hard to extend the set and deal with several ranges but my impression is that lowercase a-z is all anyone ever uses in the wild. That said, I still plan on having an extended greedy optimisation in ATen c++, maybe in a month or so.

[fritzo]

@t-vi lowercase a-z is all anyone ever uses in the wild

We're using opt_einsum.contract() for hundreds of variables in Pyro. Each local contraction requires only a few variables (hence this PR), but we have one variable per time step in a Hidden Markov Model.

[jcmgray]

Yes just to clarify, opt_einsum already maps pairwise contractions into the [a-zA-Z] range if necessary (since large contractions can have thousands of indices). The real problem is if you want 26+ indices in a single pairwise contraction - i.e. a tensor with 26+ dimensions. Sounds unusual but this is actually quite a likely/necessary situation if, for example, you are simulating large quantum circuits.

Now, opt_einsum tries to call tensordot as much as possible, so will probably avoid this niche case (torch backend + non-tensordot-able contractions + >26 dimensions), but it's certainly possible!

[t-vi]

If there is a use case, supporting A-Z is easy. One could also consider a "post-parsing interface" (i.e. taking preprocessed equations) if that is useful. I must admit I always thought of einsum as a convenient as hoc interface, so I learn something new here (thanks!).

[dgasmith]

Closed/opened to trigger Travis. Once that passes and the opt_einsum.parser.convert_to_valid_einsum_chars is either used or updated this is ready to go.

[codecov-io]

Codecov Report

Merging #41 into master will decrease coverage by 0.01%.
The diff coverage is 100%.

[fritzo]

Thanks for the quick review @dgasmith! I've moved the convert-all-chars implementation up to convert_to_valid_einsum_chars and used it in the torch backend.

[jcmgray]

Thanks for this update, all looks good to me,!

[dgasmith]

@fritzo Thanks for the PR, everything looks great and we will get this into the next release.