Use Blockwise for matmul

[ricardoV94]

We use a Blockwised Dot for matmul so that we get gradients for free.

C-performance won't be great for Blockwised Dot, since that doesn't have a C-implementation.
We could Blockwise more specialized Dot22 / GEMM Ops but that code is a bit of a mess at the moment and not useful long term as we deprecate the C backend

Alternatively we probably could use tensor_dot/batched_dot? They are fundamentally different (or it would be rather inefficient to convert one to the other)

Closes #451

Needed for pymc-devs/pymc#6897

[codecov-commenter]

Codecov Report

Merging #452 (8d4054b) into main (4efbd19) will decrease coverage by 0.02%.
Report is 6 commits behind head on main.
The diff coverage is 98.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #452      +/-   ##
==========================================
- Coverage   80.76%   80.75%   -0.02%     
==========================================
  Files         159      159              
  Lines       45869    45849      -20     
  Branches    11238    11234       -4     
==========================================
- Hits        37048    37026      -22     
- Misses       6593     6595       +2     
  Partials     2228     2228              

Files Changed	Coverage Δ
pytensor/compile/mode.py	84.40% <ø> (ø)
pytensor/ifelse.py	51.56% <ø> (ø)
pytensor/tensor/variable.py	87.42% <83.33%> (-0.34%)	⬇️
pytensor/link/jax/dispatch/scan.py	100.00% <100.00%> (ø)
pytensor/link/numba/dispatch/scan.py	95.91% <100.00%> (+0.02%)	⬆️
pytensor/scalar/basic.py	80.25% <100.00%> (ø)
pytensor/tensor/math.py	90.16% <100.00%> (-0.33%)	⬇️
pytensor/tensor/rewriting/blockwise.py	96.15% <100.00%> (+1.15%)	⬆️
pytensor/tensor/rewriting/uncanonicalize.py	96.21% <100.00%> (+0.24%)	⬆️

[jessegrabowski]

Looks great, I played around with it and everything seems to work as expected. Currently you can't compile graphs with blockwise-matmul into jax/numba, is that beyond the scope of this PR? I thought there was a jax.vectorize-type function that would make that trivial (for the jax case at least)?

[jessegrabowski]

Is all this better than a separate _matrix_vector_matmul function? I only ask because BLAS makes the distinction.

[ricardoV94]

This should be fine, once we ever go into optimizing this further in jax/numba backends we should be able to know which case is which by inspecting the input static types.

[ricardoV94]

I thought there was a jax.vectorize-type function that would make that trivial (for the jax case at least)?

Yet it should be pretty simple. It's in the todo list: #430
Interested in picking it up :) ?