PyTorch inline constants in dispatch to avoid graph breaks

[ricardoV94]

When we have static inputs, inlining helps torch not breaking the graph.

Related to #1110

---

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

[ricardoV94]

Still need to do something about the runtime broadcast in elemwise. Can we use torch._check for that instead of Python loops/asserts?

[codecov]

Codecov Report

Attention: Patch coverage is 68.88889% with 14 lines in your changes missing coverage. Please review.

Project coverage is 82.27%. Comparing base (4ea4259) to head (75eef40).
Report is 18 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/pytorch/dispatch/basic.py	52.17%	9 Missing and 2 partials ⚠️
pytensor/link/pytorch/dispatch/scalar.py	66.66%	1 Missing ⚠️
pytensor/link/pytorch/dispatch/shape.py	90.00%	1 Missing ⚠️
pytensor/link/pytorch/dispatch/subtensor.py	87.50%	1 Missing ⚠️

❌ Your patch status has failed because the patch coverage (68.88%) 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    #1118   +/-   ##
=======================================
  Coverage   82.27%   82.27%           
=======================================
  Files         186      186           
  Lines       48000    48066   +66     
  Branches     8621     8633   +12     
=======================================
+ Hits        39490    39546   +56     
- Misses       6353     6360    +7     
- Partials     2157     2160    +3     

Files with missing lines	Coverage Δ
pytensor/link/pytorch/linker.py	100.00% <100.00%> (ø)
pytensor/link/pytorch/dispatch/scalar.py	73.68% <66.66%> (-0.39%)	⬇️
pytensor/link/pytorch/dispatch/shape.py	85.71% <90.00%> (ø)
pytensor/link/pytorch/dispatch/subtensor.py	89.53% <87.50%> (-0.21%)	⬇️
pytensor/link/pytorch/dispatch/basic.py	87.40% <52.17%> (-7.10%)	⬇️

... and 13 files with indirect coverage changes

[ricardoV94]

Even without the runtime broadcast check, elemwise seems to break the graph

[Ch0ronomato]

Did you get a chance to profile this pr?

[Ch0ronomato]

Btw I did profile this. My machine actually failed to even compile dlogp for a model but I suspect that's unrelated. The logp method did show some improvement. The thing that intrigued me is this change reduced the number of guards by a lot (it was 10:1 with the other ones). I thought that maybe that was the cause of the runtime switch, but that didn't have the payoff I was expecting

[ricardoV94]

The cost of the guards may be non-linear so we should try to remove all

[Ch0ronomato]

The cost of the guards may be non-linear so we should try to remove all

Idk about removing all, since guards are the primitive that ensures runtime correctness. Significantly reduce, i agree

[Ch0ronomato]

Btw, for the actual perf benefit, these are the numbers i see.

# ricardo shape: 772 μs ± 12 μs per loop (mean ± std. dev. of 100 runs, 100 loops each)
# no ricardo shape: 818 μs ± 9.48 μs per loop (mean ± std. dev. of 100 runs, 100 loops each)

So it's like ~5%, probably more on slower cpus. The graph breaks are definitely a problem :(

[Ch0ronomato]

If we add these two flags with the changes in this PR:

torch._dynamo.config.capture_func_transforms=True
torch._dynamo.config.capture_scalar_outputs = True

we come down to almost 500us.

504 μs ± 12.2 μs per loop (mean ± std. dev. of 100 runs, 100 loops each)

[ricardoV94]

@Ch0ronomato can you revert the removal of the Elemwise bcast check (for now), and add those flags? Then we can merge this PR and keep playing with stuff

[Ch0ronomato]

The ci doesn't like those flags. I'll investigate

[Ch0ronomato]

I think the path to fix this is not use those flags by default, but when we have a shape operation. The torch compiler might be really restrictive

[Ch0ronomato]

Hopefully when #1159 gets merged we can just delete this flag altogether since torch will know these aren't dynamic