Allow overriding of transpose function

[fritzo]

Description

This PR makes backends fully overridable by checking each backed for a transpose function before using the default .transpose() method.

Before this PR, _transpose() would first look for a default implementation (x.transpose(axes)), and only if this failed look for an overridden backend transpose() function. This approach has two flaws: (1) backends cannot override the transpose function if a different .transpose() method already exists (e.g. in Pyro we want to track some metadata in transpose()); (2) the try-except may not play well with jit compilers like the PyTorch jit.

The only change in behavior is to the torch backend: before this PR torch.transpose() was used for transposes of 2 axes and torch.permute() was used for 3 or more axes; after this PR torch.permute() is always used.

Status

• Ready to go

[codecov-io]

Codecov Report

Merging #74 into master will decrease coverage by <.01%.
The diff coverage is 100%.

[jcmgray]

Yes I think this definitely makes sense. Some libraries (e.g. hptt), implement transpose directly for numpy arrays, which you would want to take precedence.

[jcmgray]

I'll merge this shall I, unless you have comments @dgasmith?

[dgasmith]

LGTM!

As per @jcmgray's comment. It would superb to get TBLIS and/or HPTT on conda-forge to try this out. Does HPTT compile to a single vectorized instruction or can it compile to multiple simultaneously like ICC multi arch or a more home grown solution as found in TBLIS?

[dgasmith]

@jcmgray Feel free to merge PR's if you wish with the usual caveat of something major should be discussed first.

[jcmgray]

As per @jcmgray's comment. It would superb to get TBLIS and/or HPTT on conda-forge to try this out. Does HPTT compile to a single vectorized instruction or can it compile to multiple simultaneously like ICC multi arch or a more home grown solution as found in TBLIS?

I'm actually not sure at all, I did add a numpy like transpose wrapper to hptt but don't know any details beyond that. Although it seems to have very good performance, the real critical transpose is the one inside numpy.tensordot it seems. On the other hand definitely agree that TBLIS would be super neat to get properly set-up.

@jcmgray Feel free to merge PR's if you wish with the usual caveat of something major should be discussed first.

Great, noted.

[fritzo]

@dgasmith would it be possible to cut a release with this PR and the new eager optimizer? I'd like to use both features in our upcoming Pyro release.

[dgasmith]

@fritzo Sure, no reason not to cut a minor release soon. Do you have a timeline that you are thinking of, might look at pulling in #73 from @jcmgray as well before the release.

[fritzo]

@dgasmith Our Pyro release is planned for the week of Dec 3. Let me know if there is any additional documentation you'd like before the next opt_einsum release. Thanks!