Improve static output shape and numba implementations of some Ops

[aseyboldt]

A collection of small fixes I have accumulated...

I thought it would be easier to just submit this as one PR instead of 5 small ones. The individual commits contain one fix each.

[ricardoV94]

Looks neat. Question is d5574c1 (fix: Fix shape of subtensors ) fixing a bug or just providing better static shape?

[aseyboldt]

Usually I just got a None shape, but in the code it seemed to just return shape==1 sometimes: d5574c1#diff-f5839e5efccae3b80e368b3d43abd60c65d64a012eb92a231c682805d4ab1e62L747

Not sure what that was about...

[ricardoV94]

Not sure what that was about...

We needed to know whether the output is broadcastable for the grad. Every thing besides that is just a nice to have that was not implemented I guess.

The reason I ask is that it makes it easier for me to review :)

[ricardoV94]

Looks good

[ricardoV94]

Tests looked like they were passing before refreshing :)

[codecov-commenter]

Codecov Report

Merging #375 (7045b91) into main (5c87d74) will decrease coverage by 0.01%.
The diff coverage is 78.98%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #375      +/-   ##
==========================================
- Coverage   80.40%   80.40%   -0.01%     
==========================================
  Files         156      156              
  Lines       45401    45500      +99     
  Branches    11106    11141      +35     
==========================================
+ Hits        36505    36582      +77     
- Misses       6689     6702      +13     
- Partials     2207     2216       +9     

Impacted Files	Coverage Δ
pytensor/link/numba/dispatch/extra_ops.py	89.53% <18.18%> (-2.71%)	⬇️
pytensor/link/numba/dispatch/basic.py	85.80% <36.84%> (-2.10%)	⬇️
pytensor/link/jax/dispatch/shape.py	91.93% <60.00%> (-2.90%)	⬇️
pytensor/scalar/basic.py	80.16% <66.66%> (-0.04%)	⬇️
pytensor/tensor/rewriting/shape.py	81.20% <83.33%> (+0.10%)	⬆️
pytensor/link/jax/dispatch/scalar.py	98.41% <100.00%> (+0.01%)	⬆️
pytensor/scan/rewriting.py	79.94% <100.00%> (ø)
pytensor/tensor/basic.py	90.77% <100.00%> (ø)
pytensor/tensor/rewriting/elemwise.py	88.99% <100.00%> (+0.49%)	⬆️
pytensor/tensor/rewriting/math.py	86.36% <100.00%> (+0.31%)	⬆️
... and 3 more

... and 1 file with indirect coverage changes

concern mentioned above

[fonnesbeck]

Can this be merged? There's a fix in here for a nutpie bug I keep running into.

[ricardoV94]

Can this be merged? There's a fix in here for a nutpie bug I keep running into.

There's a change that has to be reverted, I'll try to get on it Monday, thanks for pinging