Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cannot provide custom tensordot method #101

Closed
srush opened this issue Aug 27, 2019 · 7 comments
Closed

Cannot provide custom tensordot method #101

srush opened this issue Aug 27, 2019 · 7 comments

Comments

@srush
Copy link

srush commented Aug 27, 2019

As far as I can tell there is no way to provide a custom tensordot method. I tried pretty hard, but the code path seems to fail on examples like contract("abc,ab-> ac").

@jcmgray
Copy link
Collaborator

jcmgray commented Aug 27, 2019

Maybe you could expand a bit on what you have tried? The equation you give as an example cannot be framed as a simple tensordot, since the index 'a' appears 3 times.

@srush
Copy link
Author

srush commented Aug 28, 2019

Sure. The docs seem to imply that if you provide a tensordot and a transpose, that opteinsum does the rest. In practice though, when I tried this I ran into two issues.

  1. The code path seems to be broken, it complains that my library has no einsum function. I think there is a bug here https://github.com/dgasmith/opt_einsum/blob/master/opt_einsum/contract.py#L541

  2. When I change that line, and give it the above expression it tries to call my tensordot function and then fails with an enigmatic error (seems to get confused by the multiple instances).

In practice I love the idea that opt-einsum could be used with general operators, but it seems like maybe the docs should be updated to make it clear that it requires a full einsum implementation below the hood?

@jcmgray
Copy link
Collaborator

jcmgray commented Aug 28, 2019

The main role of opt_einsum I would consider is to find and perform pairwise contraction orderings.

With regard to point 1, does your backend provide an einsum function? If it doesn't then this isn't a bug, its just that certain pairwise contractions such as the one you give cannot be done with tensordot alone. I think the docs are reasonably clear on when an einsum implementation is required:

For a contraction to be possible without using a backend einsum, it must satisfy the following rule: in the full expression (including output indices) each index must appear twice. In other words, each dimension must be contracted with one other dimension, or left alone.

There are simple cases (such as here where 'a' is just a broadcast dimension) where one could imagine manually implementing things, but einsum has such myriad syntax uses that it gets v complex to handle, let alone performance considerations.

@srush
Copy link
Author

srush commented Aug 28, 2019

I see, I guess I got tripped up by this statement: "In fact, any library that provides a tensordot() and transpose() implementation can perform most normal contractions." Keeping axes around felt pretty normal to me (for instance any sort of batching).

Writing an einsum implementation is really hard, it requires a parser and all sort of special cases. I like the idea that opt-einsum could be used as a front-end for the non-standard contractions I want to do. Any idea for the true minimal set of functions a backend could implement such that opt-einsum could be used as a front-end?

@dgasmith
Copy link
Owner

This sounds a bit like we need a register_backend that would work something like the following:

oe.register_backend("mymodule", einsum=np.einsum, tensordot=mymodule.tensordot, transpose=mymodule.transpose)

oe.contract(..., backend="mymodule")

We should only need einsum, tensordot, and transpose for full functionality.

Another good time to look into NEP18.

@jcmgray
Copy link
Collaborator

jcmgray commented Aug 28, 2019

@dgasmith I think the issue is just the type of contraction requiring a non-existent einsum.

I see, I guess I got tripped up by this statement: "In fact, any library that provides a tensordot() and transpose() implementation can perform most normal contractions." Keeping axes around felt pretty normal to me (for instance any sort of batching).

Yes I guess 'normal' is pretty subjective! I was coming from the background of physics and einstein notation when I wrote that I suspect.

Any idea for the true minimal set of functions a backend could implement such that opt-einsum could be used as a front-end?

I think einsum can do things no other function can, but some extras that come to mind are:

  • diag (ii->i)
  • trace (ii->)
  • multiply (i,i->i)
  • sum (i->)

But once they start getting mixed together, e.g. (iij->i) it all gets a bit more complicated and complete support seems pretty impossible. Trying to support specific, common cases such as batch dimensions would probably be the way to proceed (might actually be possible with #95), but even then, there is not standard 'vectorize' function in the same way that tensordot is 'standard', so one would have to use python looping and concat or something.

@dgasmith
Copy link
Owner

dgasmith commented Nov 7, 2021

I would lean with @jcmgray that this functionality could be complex to implement. At the moment, you can hack this with the following:

import opt_einsum as oe
oe.backends.dispatch._cached_funcs["tensordor", "numpy"] = custom_func

Happy to take a PR which is well documented with edge cases.

@dgasmith dgasmith closed this as completed Nov 7, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

3 participants