Fix bug in infer_shape of Blockwise(Subtensor)

[ricardoV94]

Also adds some rewrites to make sure we end up with a clean gradient graph for a batched Conv1D.

The bug itself showed up in pymc-labs/pymc-marketing#1583

---

📚 Documentation preview 📚: https://pytensor--1353.org.readthedocs.build/en/1353/

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.04%. Comparing base (4e59f21) to head (0aeda21).
Report is 10 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1353      +/-   ##
==========================================
+ Coverage   82.02%   82.04%   +0.02%     
==========================================
  Files         203      203              
  Lines       48845    48856      +11     
  Branches     8691     8693       +2     
==========================================
+ Hits        40067    40086      +19     
+ Misses       6627     6619       -8     
  Partials     2151     2151              

Files with missing lines	Coverage Δ
pytensor/tensor/basic.py	91.17% <100.00%> (+0.03%)	⬆️
pytensor/tensor/blockwise.py	85.64% <100.00%> (ø)
pytensor/tensor/rewriting/blockwise.py	96.68% <100.00%> (+0.28%)	⬆️

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jessegrabowski]

Not sure I get the objective, but the code looks good

[jessegrabowski]

Why does copying the input simulate an operation on the inputs? The apply takes the copies as inputs, so any intermediate operations (the f in a_identity = f(a)) would be lost right?

[ricardoV94]

After rewrites, but the Blockwise.infer_shape acts on the returned graph immediately.

The thing Blockwise is trying to is to figure out if the the core_shape of the Op depends on the values of the inputs, or can be guessed just based on their shapes. For that it calls infer_shape on the core op with dummy core variables and then checks if those are part of the returned shape graph. If they are, then it can't really use it, because it means the shape may vary over iterations (say a Blockwise slice Subtensor with batch start points).

If they are not used in the shape of the graph, then it means only the core shape is needed which is fine to use. This would be the case of a Blockwise(Dirichlet), where only the core shape (length of alpha), but not their value is needed. (We don't blockwise RVs but you get the idea).

Anyway, the logic to figure out if the core values are needed was to create a dummy node and then checking if the node inputs were in the graph. But this failed when the dummy node didn't really use the dummy variables (because it added extra nodes, like DimShuffle or ScalarFromTensor in the case that actually failed). Identity here is an easy way to test this without having to change anything else in the test.