Add `Blockwise` `Op`

[purna135]

This PR builds off of #757 and closes #695.

To #757 it adds:

• get_output_info(), which is the same as Elemwise get_output_info(), to make all inputs of the same dimension.
• derive DimShuffle's gufunc signature
• reduce the broadcasted dimensions of inputs after the grad is computed

Differences with #757:

• instead of using the dimensions from the start for computing the curr_static_shape of core_inp_grads use the dimensions from the end.
• an extra check before calling perform() of DimShuffle (which can be removed later)

[codecov[bot]]

Codecov Report

Merging #1215 (6cda5c3) into main (462d8d5) will increase coverage by 4.14%. The diff coverage is 86.00%.

:exclamation: Current head 6cda5c3 differs from pull request most recent head c7b0d10. Consider uploading reports for the commit c7b0d10 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1215      +/-   ##
==========================================
+ Coverage   75.02%   79.16%   +4.14%     
==========================================
  Files         194      174      -20     
  Lines       50099    48677    -1422     
  Branches    12096    10359    -1737     
==========================================
+ Hits        37586    38536     +950     
+ Misses      10189     7640    -2549     
- Partials     2324     2501     +177     

Impacted Files	Coverage Δ
aesara/tensor/blockwise.py	85.81% <85.81%> (ø)
aesara/tensor/math.py	90.05% <100.00%> (-0.62%)	:arrow_down:

... and 122 files with indirect coverage changes

[brandonwillard]

Don't forget to rebase onto upstream/main (or whatever the name of your remote for this repository is); that should remove the merge commit in this branch.

[purna135]

What should be the signature for Subtensor Op and Shape Op ?

[brandonwillard]

What should be the signature for Subtensor Op and Shape Op ?

If you're talking about constructing symbolic graphs, the signatures are ultimately determined by their Op.make_node implementations.

[purna135]

If you're talking about constructing symbolic graphs, the signatures are ultimately determined by their Op.make_node implementations.

Yes, got it now

[purna135]

Hello, @brandonwillard. I'm having some DimShuffle related problems that I can't figure out. Could you please take a look and assist in determining which piece of logic is causing this error?

You can reproduce the error using the following command. pytest tests/tensor/test_blockwise.py::test_blockwise_solve_grad[a_shape0-b_shape0]

[brandonwillard]

Hello, @brandonwillard. I'm having some DimShuffle related problems that I can't figure out. Could you please take a look and assist in determining which piece of logic is causing this error?

You can reproduce the error using the following command. pytest tests/tensor/test_blockwise.py::test_blockwise_solve_grad[a_shape0-b_shape0]

It looks like SolveBase.L_op is producing tensors with an extra dimension that Solve2 can't handle.

I'm guessing Solve2 was meant to serve as a specialization of Solve's more general (matrix x matrix) -> matrix signature, but its inherited L_op probably doesn't match the signature change.

Regardless, we shouldn't need new Ops for that; instead, a helper function like aesara.tensor.slinalg.solve can be used to project the inputs and outputs to and from Solve's signature's space.