Add torch.tensordot to the PyTorch converter - #2846
Open
rohanpoudel2 wants to merge 2 commits into
Open
Conversation
torch.tensordot reaches the converter as a single aten::tensordot node under both TorchScript and torch.export, so it needs its own lowering. Contract it by gathering the contracted dimensions together and folding the contraction into one matmul, then restoring the free dimensions. Fixes apple#2137
aten::tensordot is not in the Core ATen opset, so to_edge decomposes it into permute_copy, _clone_dim_order, view_copy and mm. The converter is never reached on that path, and dim_order_ops::_clone_dim_order has no converter of its own, so exercising it here only reports an unrelated gap.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #2137.
torch.tensordothas no PyTorch-to-MIL converter, so models that use it fail with:tensordotis not decomposed by either frontend — it reaches the converter as asingle
aten::tensordotnode under TorchScript, and survivesrun_decompositions({})asaten.tensordot.defaultundertorch.export— so itneeds its own lowering. This adds one for both.
Implementation
A
tensordotis a permutation away from a plain matrix multiplication, so thelowering gathers the contracted dimensions together and folds the whole
contraction into a single
matmul:selfto(free dims..., contracted dims...)andotherto(contracted dims..., free dims...), reshape both to rank 2,matmul, thenrestore the free dimensions.
matrix multiply such as
dims=([1], [0])lowers to a singlematmulop ratherthan to no-op transposes for the MIL passes to clean up afterwards.
-1, so a free dimension that is onlyknown at run time is preserved. When one is symbolic the output shape is read
back off the inputs with
gatheronshapeinstead of being baked in.torch.tensordot.of both inputs (a scalar result) both fall out of the same path.
mismatched contraction raises a
ValueErrornaming the two dimensions.Tests
Adds
TestTensordotcovering, across TorchScript andtorch.export:every dimension of the smaller input, non-trailing contracted dimensions, and
inputs of differing rank;
66 passed on macOS (arm64), across
mlprogramfp16 andneuralnetworkfp32, onthe TorchScript and
torch.exportfrontends. The neighbouringTestCdist,TestOuter,TestDot,TestMvandTestMatMulclasses still pass unchanged.The ExecuTorch frontend is skipped.
aten::tensordotis not in the Core ATenopset, so
to_edgedecomposes it intopermute_copy,_clone_dim_order,view_copyandmm; this converter is never reached on that path, anddim_order_ops::_clone_dim_orderhas no converter of its own, so running thetests there only reports an unrelated gap. Happy to drop the skip if you would
rather see that gap tracked from here.
The lowering was also checked directly against
torch.tensordoton randomizedrank-1 to rank-4 inputs with randomized contracted axes, on both supported
frontends: 44/44 matched in shape and value, with the remaining cases skipped
only because their outer-product result exceeds Core ML's rank-5 limit. The
reproduction in the issue converts and matches PyTorch to 2.4e-07.