Implement graph.vectorize and Blockwise Op

[ricardoV94]

Critical for pymc-devs/pymc#5383

This PR was done almost completely from scratch, although the general principles and some of the code were obviously inspired by aesara-devs/aesara#1215

For fairness, all authors of the commits in that PR were added as co-authors here.

CC @purna135

Some design questions:

1. Do we want to make Blockwise the "default" form of an Op (in terms of rewrites), and only later specialize if they were not needed? Decided yes
2. Do we want to fuse consecutive Blockwises with the same batched dims, similarly to how we fuse consecutive Elemwises? Not for now
3. Can Rop be implemented just as the batched core Rop, like L_op?

Important follow-up:

1. Dispatch to JAX and Numba
2. Create a PyTensor "vectorize" function that vectorizes the whole graph, similar to what the L_op method has to do Included as the last commit

from pytensor.graph import vectorize

# Original graph
x = pt.vector("x")
y = pt.exp(x) / pt.sum(pt.exp(x))

new_x = pt.matrix("new_x")
[new_y] = vectorize([y], {x: new_x})

fn = pytensor.function([new_x], new_y)
fn([[0, 1, 2], [2, 1, 0]])
# array([[0.09003057, 0.24472847, 0.66524096],
#        [0.66524096, 0.24472847, 0.09003057]])

Important to test:

1. Whether this works with Ops with inner graphs (Scan, Composite)

[codecov-commenter]

Codecov Report

Merging #306 (ff1f7c5) into main (8112576) will increase coverage by 0.01%.
The diff coverage is 81.34%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #306      +/-   ##
==========================================
+ Coverage   80.39%   80.40%   +0.01%     
==========================================
  Files         156      158       +2     
  Lines       45397    45652     +255     
  Branches    11106    11181      +75     
==========================================
+ Hits        36497    36708     +211     
- Misses       6694     6722      +28     
- Partials     2206     2222      +16     

Impacted Files	Coverage Δ
pytensor/link/numba/dispatch/nlinalg.py	100.00% <ø> (ø)
pytensor/tensor/basic.py	90.77% <ø> (+0.01%)	⬆️
pytensor/tensor/rewriting/linalg.py	76.92% <64.91%> (+4.70%)	⬆️
pytensor/tensor/blockwise.py	79.03% <79.03%> (ø)
pytensor/tensor/slinalg.py	92.30% <88.63%> (-1.25%)	⬇️
pytensor/tensor/elemwise.py	88.44% <92.00%> (+0.42%)	⬆️
pytensor/compile/mode.py	84.47% <100.00%> (ø)
pytensor/tensor/nlinalg.py	98.07% <100.00%> (+0.23%)	⬆️
pytensor/tensor/random/op.py	97.51% <100.00%> (+0.04%)	⬆️
pytensor/tensor/rewriting/blockwise.py	100.00% <100.00%> (ø)

... and 7 files with indirect coverage changes

[ferrine]

I see issues with the design approach. The rewrites become more and more involved and should take blockwise and non blockwise ops in account. I think a discussion over this issue is required before any steps forward.

[ricardoV94]

@ferrine I think the sanest thing is for rewrites to only target Blockwise graphs, the same way that our rewrites only target Elemwise graphs and never the equivalent Scalar graphs (see some discussion in #107).

Users would now always work with Blockwise (that's what we expose by default when you do pt.foo), the same way they always work with Elemwise. During compilation we remove Blockwise when unnecessary the same way that #107 suggests to remove Elemwise when unnecessary, but only near the end of compilation.

What problem do you see with this?

[ferrine]

Let's get more reviews

[ricardoV94]

Last commits introduces a vectorize utility. This should be the PyTensor equivalent to np.vectorize or jax.vectorize.

If people disagree too much with the API I will drop this commit so we can discuss it elsewhere. I would like to get the Blockwise functionality merged regardless of what the final vectorize may look like, because this functionality is sorely missing.

See new example in top comment

[twiecki]

I think in pymc4 we tried to use functionality like this to sample multiple chains in parallel (in TF). Sounds like we could just take a pymc model logp and vectorize it along the dim of chains. Now the question is whether that would provide any benefits.

[ricardoV94]

I think in pymc4 we tried to use functionality like this to sample multiple chains in parallel (in TF). Sounds like we could just take a pymc model logp and vectorize it along the dim of chains. Now the question is whether that would provide any benefits.

I am not sure what numpyro vectorize does, but that sounds a bit strange? The sampler probably wants to go different tree_depths for different chains / points?

I see this as being more useful for actual model building, where you define the core case and then vectorize batched dims.

Also useful for marginalization of discrete variables via enumeration that pymc_experimental.marginal_model does.

Might also be useful for stuff like compute_log_likelihood and prior and posterior predictive sampling where we now resort to iterating over chains/draws, although I am not sure what cases would actually be faster.

[twiecki]

I am not sure what numpyro vectorize does, but that sounds a bit strange? The sampler probably wants to go different tree_depths for different chains / points?

Right, for NUTS it's a pain, as @ColCarroll can attest to. I think that's where some of the new sampler developments might come in handy by Matt Hoffman et al.

I see this as being more useful for actual model building, where you define the core case and then vectorize batched dims.

Also useful for marginalization of discrete variables via enumeration that pymc_experimental.marginal_model does.

Might also be useful for stuff like compute_log_likelihood and prior and posterior predictive sampling where we now resort to iterating over chains/draws, although I am not sure what cases would actually be faster.

Interesting potential applications.

In any case, I think the API of vectorize is so simple and elegant that I find it hard to imagine someone to take issue with it. It's also very analogous to the existing numpy functionality.

[ricardoV94]

Tests are passing again

[twiecki]

Merge?

[ricardoV94]

Merge?

I think so. Nothing like trying it out there to see if the design is good