Update Scan optimizations and fix Scan + RandomVariable shape inference issues

[brandonwillard]

This PR refactors some of the Scan optimizations in order to fix Scan + RandomVariable shape inference issues. These changes are a part of #584 and are a solution to #608.

• Replace use of FunctionGraph.replace* in add_nitsot_outputs
  This isn't directly related to Fix PushOutNonSeqScan + RandomVariable shape inference error in Scans #608, so it can be taken care of later.

• Use SpecifyShape with the RandomVariable's size parameter
  This makes it possible for Scan's shape inference to track the shape of the cloned non-sequence variables it makes out of RandomVariable size parameters.

• Add canonicalizations for SpecifyShape
  These are necessary for broadcastable inference in RandomVariable when SpecifyShape is used.

• Convert global Scan optimizations to local optimizations.

  It looks like the current Scan optimizations mutate the FunctionGraphs in-place in different stages, and this might be leading to inconsistent intermediate graph states that fail when the shape inference optimization is performed.

  This seems to be the case for the MWE in Fix PushOutNonSeqScan + RandomVariable shape inference error in Scans #608, because shape inference errors occur when the cloned size parameter in the Scan's body (i.e. the size variable in the graph constructed by the scan_body function) is encountered. It should be possible to assign the shape of size_at to that variable in the ShapeFeature; however, the outer size_at is nowhere to be found in the FunctionGraph that's being optimized—implying that a valid Scan isn't being used to replace the original Scan, but that the original Scan is possibly being changed in multiple FunctionGraph.replace* steps.

  From inspection of aesara.graph.opt, it's clear that this is a distinct possibility, since most of the optimizations are global optimizations with multiple FunctionGraph.replace_all* calls in a single optimization. This approach to graph rewriting is rather undesirable for the reasons mentioned above, so this PR will attempt to remedy the situation.

[codecov]

Codecov Report

Merging #635 (afefb01) into main (cbf9112) will decrease coverage by 0.01%.
The diff coverage is 86.44%.

@@            Coverage Diff             @@
##             main     #635      +/-   ##
==========================================
- Coverage   77.16%   77.14%   -0.02%     
==========================================
  Files         156      156              
  Lines       47007    47022      +15     
  Branches    10281    10282       +1     
==========================================
+ Hits        36271    36277       +6     
- Misses       8150     8157       +7     
- Partials     2586     2588       +2     

Impacted Files	Coverage Δ
aesara/compile/ops.py	83.57% <ø> (-0.12%)	⬇️
aesara/graph/opt.py	63.21% <0.00%> (ø)
aesara/scan/op.py	81.98% <ø> (ø)
aesara/scan/opt.py	81.30% <ø> (-0.75%)	⬇️
aesara/tensor/type.py	91.21% <ø> (-0.03%)	⬇️
aesara/scan/utils.py	87.25% <50.00%> (ø)
aesara/tensor/shape.py	88.40% <60.00%> (-0.42%)	⬇️
aesara/tensor/basic_opt.py	84.67% <77.77%> (-0.04%)	⬇️
aesara/tensor/subtensor_opt.py	84.33% <91.30%> (+0.16%)	⬆️
aesara/link/numba/dispatch/random.py	100.00% <100.00%> (ø)
... and 6 more