Skip to content

Add torch.tensordot to the PyTorch converter - #2846

Open
rohanpoudel2 wants to merge 2 commits into
apple:mainfrom
rohanpoudel2:tensordot-support
Open

Add torch.tensordot to the PyTorch converter#2846
rohanpoudel2 wants to merge 2 commits into
apple:mainfrom
rohanpoudel2:tensordot-support

Conversation

@rohanpoudel2

@rohanpoudel2 rohanpoudel2 commented Sep 1, 2026

Copy link
Copy Markdown

Summary

Fixes #2137.

torch.tensordot has no PyTorch-to-MIL converter, so models that use it fail with:

RuntimeError: PyTorch convert function for op 'tensordot' not implemented.

tensordot is not decomposed by either frontend — it reaches the converter as a
single aten::tensordot node under TorchScript, and survives
run_decompositions({}) as aten.tensordot.default under torch.export — so it
needs its own lowering. This adds one for both.

Implementation

A tensordot is a permutation away from a plain matrix multiplication, so the
lowering gathers the contracted dimensions together and folds the whole
contraction into a single matmul:

  • Permute self to (free dims..., contracted dims...) and other to
    (contracted dims..., free dims...), reshape both to rank 2, matmul, then
    restore the free dimensions.
  • The permutes are skipped when they are already the identity, so an ordinary
    matrix multiply such as dims=([1], [0]) lowers to a single matmul op rather
    than to no-op transposes for the MIL passes to clean up afterwards.
  • The free dimensions are reshaped with -1, so a free dimension that is only
    known at run time is preserved. When one is symbolic the output shape is read
    back off the inputs with gather on shape instead of being baked in.
  • Negative dimensions are wrapped, matching torch.tensordot.
  • Contracting no dimensions (the outer product) and contracting every dimension
    of both inputs (a scalar result) both fall out of the same path.
  • Contracted dimensions must be static and must agree pairwise; a symbolic or
    mismatched contraction raises a ValueError naming the two dimensions.

Tests

Adds TestTensordot covering, across TorchScript and torch.export:

  • the shapes from the issue, plus a single contracted dimension, contracting
    every dimension of the smaller input, non-trailing contracted dimensions, and
    inputs of differing rank;
  • negative dimensions;
  • no contraction at all (outer product);
  • full contraction of both inputs (scalar result);
  • a free dimension that is only known at run time.
python -m pytest coremltools/converters/mil/frontend/torch/test/test_torch_ops.py -k TestTensordot -q

66 passed on macOS (arm64), across mlprogram fp16 and neuralnetwork fp32, on
the TorchScript and torch.export frontends. The neighbouring TestCdist,
TestOuter, TestDot, TestMv and TestMatMul classes still pass unchanged.

The ExecuTorch frontend is skipped. aten::tensordot is not in the Core ATen
opset, so to_edge decomposes it into permute_copy, _clone_dim_order,
view_copy and mm; this converter is never reached on that path, and
dim_order_ops::_clone_dim_order has no converter of its own, so running the
tests 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.tensordot on randomized
rank-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.

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.
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

Successfully merging this pull request may close these issues.

Runtime error, op 'tensordot' not implemented.

1 participant