Add shape_unsafe tag to rewrites that can hide shape errors

[ricardoV94]

• Remove BroadcastTo in favor of Alloc
• Address spoken inconsistencies between Second / Alloc in rewrites
• Add tag to easily exclude rewrites that can hide shape errors
• Simplify such rewrites
• Improve static shape of Alloc

Note that from a user standpoint, providing static shapes (via vector("x", shape=(5,)) or specify_shapes) will many times reveal shape errors immediately (this is the case for 99% of PyMC models). In this case users should feel pretty safe about "shape_unsafe" rewrites because they aren't really masking anything that wasn't checked before already.

Alloc.make_node now also raises early when it can see the provided shape is inconsistent. Alloc and Elemwise make up all of the tagged "shape_unsafe" rewrites so far.

With this PR, users can also do mode=get_default_mode().excluding("shape_unsafe") or add shape_unsafe to the excluding config to skip these rewrites at the cost of less optimizations.

Closes #367

[ricardoV94]

We discussed with @aseyboldt that it may make sense to generalize Second so that it accepts arbitrary many inputs and returns every variable as output. This would become a flat broadcast_arrays once Elemwised, and make rewrites easier to read. By overriding the __str__ we can also make it much more readable in debug_print than the current nested Second

[codecov-commenter]

Codecov Report

Merging #381 (2a3adbe) into main (7218431) will decrease coverage by 0.03%.
The diff coverage is 90.90%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #381      +/-   ##
==========================================
- Coverage   80.44%   80.41%   -0.03%     
==========================================
  Files         156      156              
  Lines       45470    45413      -57     
  Branches    11136    11119      -17     
==========================================
- Hits        36578    36520      -58     
- Misses       6687     6693       +6     
+ Partials     2205     2200       -5     

Impacted Files	Coverage Δ
pytensor/link/jax/dispatch/extra_ops.py	72.58% <ø> (-3.48%)	⬇️
pytensor/link/numba/dispatch/extra_ops.py	91.78% <ø> (-0.47%)	⬇️
pytensor/tensor/rewriting/elemwise.py	88.99% <ø> (ø)
pytensor/tensor/rewriting/math.py	86.32% <72.72%> (-0.05%)	⬇️
pytensor/tensor/rewriting/basic.py	94.01% <95.74%> (+0.48%)	⬆️
pytensor/configdefaults.py	65.92% <100.00%> (-0.10%)	⬇️
pytensor/tensor/basic.py	90.82% <100.00%> (+0.05%)	⬆️
pytensor/tensor/extra_ops.py	88.53% <100.00%> (-0.48%)	⬇️
pytensor/tensor/rewriting/extra_ops.py	88.23% <100.00%> (-1.02%)	⬇️

... and 9 files with indirect coverage changes

[ricardoV94]

Fixing #379 should also help with the "unsafety" concerns

[aseyboldt]

BroadcastTo is imported in pymc a couple of times. Maybe we should leave an empty Op here, that is deprecated and doesn't do anything?

[ricardoV94]

Some of the removed rewrites are also directly imported.

This shouldn't be a problem however. I marked this PR as a major release so we will bump the version above the upper-bound pinned by PyMC. When we update the pin on PyMC I'll address the changes. They require some manual review anyway to see if the logic that depended on BroadcastTo was valid per our new rules and can be transferred to Alloc.

This was all on the logprob inference module AFAICT so impact should be pretty contained.

[aseyboldt]

sounds good

[aseyboldt]

Other than the two suggestions above this looks good :-)