Implement several subtensor lift rewrites

[ricardoV94]

This allows reducing computations on batch dimensions by lifting simple indexing operations closer to the inputs.

An obvious example is:

import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import get_default_mode

mode = get_default_mode()

x = pt.matrix("x", shape=(512, 512))
x_test = np.random.normal(size=x.type.shape)

x_sum = x.sum(axis=1)
out = x_sum[0]

fn_before = pytensor.function([x], out, mode=mode.excluding("local_subtensor_of_reduce"))
fn_before.dprint(print_type=True)
%timeit fn_before(x_test)
# Subtensor{i} [id A] <Scalar(float64, shape=())> 1
#  ├─ Sum{axis=1} [id B] <Vector(float64, shape=(512,))> 0
#  │  └─ x [id C] <Matrix(float64, shape=(512, 512))>
#  └─ 0 [id D] <uint8>
# 762 μs ± 7.55 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

print()
fn_after = pytensor.function([x], out, mode=mode)
fn_after.dprint(print_type=True)
%timeit fn_after(x_test)
# Sum{axes=None} [id A] <Scalar(float64, shape=())> 1
#  └─ Subtensor{i} [id B] <Vector(float64, shape=(512,))> 0
#     ├─ x [id C] <Matrix(float64, shape=(512, 512))>
#     └─ 0 [id D] <uint8>
# 5.26 μs ± 86 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

---

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

[codecov]

Codecov Report

Attention: Patch coverage is 90.97561% with 37 lines in your changes missing coverage. Please review.

Project coverage is 82.07%. Comparing base (5335a68) to head (4b545cf).
Report is 12 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/rewriting/subtensor_lift.py	90.77%	17 Missing and 20 partials ⚠️

❌ Your patch status has failed because the patch coverage (90.97%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1158      +/-   ##
==========================================
+ Coverage   82.02%   82.07%   +0.05%     
==========================================
  Files         207      208       +1     
  Lines       49301    49517     +216     
  Branches     8747     8785      +38     
==========================================
+ Hits        40440    40642     +202     
- Misses       6695     6702       +7     
- Partials     2166     2173       +7     

Files with missing lines	Coverage Δ
pytensor/tensor/rewriting/subtensor.py	90.08% <100.00%> (+0.49%)	⬆️
pytensor/tensor/rewriting/subtensor_lift.py	90.77% <90.77%> (ø)

🚀 New features to boost your workflow:

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

[lucianopaz]

I finally got around to reviewing this. Sorry @ricardoV94 for the huge delay! I left a few questions but it looks good.

[lucianopaz]

I don't know the difference between canonicalize and specialize.

[ricardoV94]

canonicalize happens before stabilize and before specialize. Canonicalize should simplify graphs convert different forms into a canonical one. Specialize is allowed to go bananas and produce very specialized code. Many rewrites we apply in multiple phases so they can interact with others that are specific to that phase

[lucianopaz]

Will the elemenwise inputs already have the same number of dimensions as the output at this point?

[ricardoV94]

Yes, Elemwise.make_node adds any expand_dims needed so that all inputs/output have the same ndim

[lucianopaz]

What does this function do?

[ricardoV94]

When there's a runtime error evaluating a node it shows the stack trace of where the original variable was first defined by the user, even if it was later replaced by rewrites and no longer really exists in the computational graph

[lucianopaz]

Does this rewrite handle the case where keepdims=True? Or does that kwarg produce a different Op than CAReduce?

[ricardoV94]

keepdims is not part of CARedue, it's added as a separate expand_dims after the reduction.

[lucianopaz]

Add pragma: no cover to make codecov shut up

[lucianopaz]

This rewrite made me question something of the other commits. If this rewrite is lifting the subtensor out of the advanced subtensor operation, aren't some of your other rewrites actually lowering the subtensor so that it gets applied before the other ops?

[ricardoV94]

when we say lifting we mean applying closer to the inputs of the function (before other operations). What you're saying would be sinking/lowering. In general we want to apply the subtensor before other elementwise operations, so as to compute fewer things). AdvancedIndexing could be a case to lower, if we end up having more rows than before, but we don't have any rewrites to reason about that