Rewrite batched dots that do not reduce as multiplication

[ricardoV94]

Dot is basically a fused broadcasted multiplication and reduction. When we have cases that correspond just to multiplication there is no advantage of using Dot, specially in vectorized graphs where we'll have Blockwised versions of the Dot.

This PR adds a rewrite for these cases.

Initially I was rewritting the non-blockwise versions as well, but those interfere with the BLAS pipeline. Testing locally BLAS for the outer product is not always faster (depends on the size), but it's generally faster for when doing the rank1 matrix update, so I don't want to mess with that for now.

I think we should introduce BLAS in those cases only when doing an update, not when doing the outer product alone. The rewrite is easy to toggle to allow addressing the non-blockwised version.

---

📚 Documentation preview 📚: https://pytensor--1178.org.readthedocs.build/en/1178/

[codecov]

Codecov Report

Attention: Patch coverage is 85.71429% with 7 lines in your changes missing coverage. Please review.

Project coverage is 82.27%. Comparing base (b065112) to head (8e08d2f).
Report is 6 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/rewriting/math.py	86.66%	2 Missing and 2 partials ⚠️
pytensor/tensor/math.py	84.21%	2 Missing and 1 partial ⚠️

❌ Your patch status has failed because the patch coverage (85.71%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1178   +/-   ##
=======================================
  Coverage   82.27%   82.27%           
=======================================
  Files         186      186           
  Lines       48009    48053   +44     
  Branches     8624     8633    +9     
=======================================
+ Hits        39499    39538   +39     
- Misses       6353     6356    +3     
- Partials     2157     2159    +2     

Files with missing lines	Coverage Δ
pytensor/tensor/math.py	91.87% <84.21%> (+0.01%)	⬆️
pytensor/tensor/rewriting/math.py	88.93% <86.66%> (-0.05%)	⬇️

[jessegrabowski]

Looks nice. Added some nitpicks, nothing major. Have you benchmarked this?

[jessegrabowski]

it's match week here at pymc-devs!

[jessegrabowski]

Why did you refactor check_alloc_runtime_broadcast out of this TestAlloc class but not this one?

[ricardoV94]

I refactored the other out because trying to use the method from another module led pytest to execute all the tests of the class. The class calling a method on the class itself should be fine, which I assume is what's happening here?

[ricardoV94]

Looks nice. Added some nitpicks, nothing major. Have you benchmarked this?

I'm sure it's faster without benchmarking because blockwise will be a python loop over batch dims. Even the core case without blockwise was often faster than with dot -> GER (except n=512 haha), but I don't want to mess with the rest of the BLAS pipeline as I mentioned