Exclude backend incompatible rewrites in Scan dispatch

[ricardoV94]

The whole inner mode is still a problem, but for now another bandaid

Closes #426

[codecov-commenter]

Codecov Report

Merging #427 (ea9f0f4) into main (6d3c756) will decrease coverage by 0.03%.
Report is 7 commits behind head on main.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #427      +/-   ##
==========================================
- Coverage   80.78%   80.75%   -0.03%     
==========================================
  Files         157      157              
  Lines       45575    45612      +37     
  Branches    11152    11162      +10     
==========================================
+ Hits        36816    36836      +20     
- Misses       6552     6562      +10     
- Partials     2207     2214       +7     

Files Changed	Coverage Δ
pytensor/link/jax/dispatch/scan.py	100.00% <100.00%> (ø)
pytensor/link/numba/dispatch/scan.py	95.91% <100.00%> (+0.02%)	⬆️

... and 4 files with indirect coverage changes

[jessegrabowski]

Let me speed test the default mode without rewrites before this gets merged

[ricardoV94]

Nudge @jessegrabowski

[jessegrabowski]

This is good to go. The speed tests I ran show excluding the BLAS rewrites has very little effect, even on large-scale problems:

Full notebook here. I'm actually a bit disappointed by these results. I thought the rewrites to the lower-level BLAS operations like GEMM would offer huge speedups relative to whatever the default is, but it seems like the gains are quite marginal. Maybe whatever routine the dot Op is calling in C also uses these same routines?

[ricardoV94]

@jessegrabowski the GEMM/BLAS OPs are only relevant for the C backend. Jax and Numba introduce it themselves when appropriate, so the changes shouldn't have any effect.

We didn't change the rewrites used in the C backend

[jessegrabowski]

OK, I misunderstood the PR, I thought you were globally disabling JAX incompatible rewrites for the inner graphs of scan. Sorry for holding this up.

[ricardoV94]

NP