Implement vectorize_node for CheckAndRaise Op

[ricardoV94]

Description

Implement vectorize dispatch for Assert. This crops up a lot when vectorizing logp graphs in PyMC.

Related Issue

• Related to Vectorize follow-up #430

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[codecov-commenter]

Codecov Report

Merging #553 (bb3fda4) into main (c38eea0) will increase coverage by 0.00%.
Report is 2 commits behind head on main.
The diff coverage is 94.44%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #553   +/-   ##
=======================================
  Coverage   80.90%   80.91%           
=======================================
  Files         162      162           
  Lines       46416    46428   +12     
  Branches    11354    11356    +2     
=======================================
+ Hits        37554    37568   +14     
+ Misses       6639     6638    -1     
+ Partials     2223     2222    -1     

Files	Coverage Δ
pytensor/breakpoint.py	55.00% <ø> (ø)
pytensor/compile/builders.py	77.17% <100.00%> (ø)
pytensor/graph/basic.py	89.08% <ø> (ø)
pytensor/graph/rewriting/kanren.py	100.00% <ø> (ø)
pytensor/ifelse.py	51.70% <100.00%> (ø)
pytensor/link/jax/dispatch/random.py	96.05% <100.00%> (ø)
pytensor/link/jax/dispatch/shape.py	87.09% <ø> (ø)
pytensor/link/jax/dispatch/subtensor.py	88.13% <ø> (ø)
pytensor/link/jax/dispatch/tensor_basic.py	91.96% <ø> (ø)
pytensor/link/numba/dispatch/random.py	98.83% <100.00%> (ø)
... and 36 more

[jessegrabowski]

Why does vectorizing y make a second, but vectorizing x doesn't? second is the strange dummy ExpandDims (Dimshuffle?) Op that always gets rewritten away right?

[ricardoV94]

y is never an output of CheckAndRaise, so it doesn't need to be vectorized (and it can't, since the Op demands it be a scalar bool). The condition is either a scalar True or False, and broadcasting (and then reducing again with all) would be a no-op.

Second is the Elemwise implicit form of Broadcasting, that allows easier rewrites during canonicalization. When needed in the final graph it gets eventually rewritten as an Alloc:

pytensor/pytensor/tensor/rewriting/basic.py

Lines 1 to 23 in 9df55cc

	""" Tensor optimizations addressing the ops in basic.py.
	Notes
	-----
	There are two ways of broadcasting arrays:
	second(x, y) == alloc(y, broadcast_shapes(x.shape, y.shape))
	The second can be more efficient because x doesn't usually need to be computed when we only want its shape.
	It may also allow other rewrites that don't try to modify x when it has multiple clients (for fear of duplicating computation).
	However, the first one is easier to reason about.
	Knowing we have such a graph allows to do certain rewrites such as "sinking" broadcasting operations below Elemwise.
	The same rewrites with alloc would be more complicated as we would need to symbolically combine the shapes of each one.
	As an example contrast rewriting the following two equivalent graphs
	alloc(x, broadcast_shapes(x.shape, y.shape)) + alloc(y, broadcast_shapes(x.shape, y.shape)) -> x + y
	second(y, x) + second(x, y) -> x + y
	Theano developers (mostly) preferred to use the first form during canonicalization and introduce the second form later,
	via rewrites like `local_fill_to_alloc`, and using the `alloc_like` helper inside rewrites.
	Many stabilize and stabilization rewrites refuse to be applied when a variable has multiple clients, so this is important.
	"""