Fix tensordot implementation

[lucianopaz]

Description

This PR discards the past implementation of tensordot and uses the one from numpy. The past implementation had many problems:

• It had two completely independent branches of execution so code was not thoroughly testes.
• It was recursive. If you passed in axes as a sequence of sequences, it tried to do transpose the dimensions and then call tensordot again using axes as an integer. This made things harder to maintain.
• It had bugs as shown by BUG: tensordot is broken #606.
• It lost the static shape information.

Taking the implementation from numpy works well. It handles all the cases of axes, it has a single execution branch, and preserves the static shape information. The only downside is that it does not have the logic to implement batched_tensordot. From my point of view, this is a good thing. batched_dot and batched_tensordot were poor-men implementations of what we can now do using Blockwise and better vectorization. I think that those two could be deprecated altogether and eventually removed.

I still haven't added tests for this PR because I want to see the current test suite coverage report. The coverage report showed that there was a test that needed to be adapted, but most of tensordot was already covered by the existing test suite. I added special tests for #606, static shape and runtime shape validation.

Related Issue

• Closes BUG: tensordot is broken #606
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (f799219) 80.87% compared to head (b7402d6) 80.86%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #607      +/-   ##
==========================================
- Coverage   80.87%   80.86%   -0.01%     
==========================================
  Files         162      162              
  Lines       46680    46743      +63     
  Branches    11408    11419      +11     
==========================================
+ Hits        37751    37800      +49     
- Misses       6699     6705       +6     
- Partials     2230     2238       +8     

Files	Coverage Δ
pytensor/tensor/math.py	89.49% <100.00%> (-0.57%)	⬇️

... and 1 file with indirect coverage changes

[ricardoV94]

Btw this looks much better than the previous implementation!

We should open an issue to deprecate the old tensordot as dot utility

[lucianopaz]

Btw this looks much better than the previous implementation!

We should open an issue to deprecate the old tensordot as dot utility

I'll open an issue for it

[ricardoV94]

Some small docstring tweak and a possible nitpick in the test (feel free to ignore) and this is good from my side

[ricardoV94]

Thanks @lucianopaz !