[WIP] Allow function dispatch for constants

[Ch0ronomato]

Description

Pytorch doesn't seem to handle the constants we put in the globals of the generated code well. It can cause graph breaks for things that shouldn't be graph break (example: The index in IncSubtensor when the index is static). This change allows the typify method to instead return a function, and then defer fetching the constants in the generated code as a function call.

A small example:

import pytensor
import pytensor.tensor as ptt

x = ptt.vector('x')
o = x[-1].inc(1)
f = pytensor.function(inputs=[x], outputs=o, mode="PYTORCH")

will generate

def pytorch_funcified_fgraph(x):
    tensor_constant = torch_assume_constant()
    scalar_constant = torch_assume_constant_1()
    # IncSubtensor{i}(x, 1, -1)
    tensor_variable = inc_subtensor(x, tensor_constant, scalar_constant)
    return (tensor_variable,)

torch_assume_constant can now do whatever it needs to do to allow the backend (torch in this case) to know that the variable is infact a constant.

Related Issue

• Closes IncSubtensor causes graph break in pytorch backend #1154
• Related to #

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):

---

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

[ricardoV94]

So did it work?

[ricardoV94]

This seems wrong?

[Ch0ronomato]

Say more?

The torch compiler threw asserts when a zero dim np value was passed back.

[ricardoV94]

Isn't it going to upcast values to float64, whatever python integers are? Does torch have scalars (not tensors) with specific dtypes we can use instead?

[Ch0ronomato]

It worked for the very simple case. I'm gonna try it with the 8 schools model and see if that works.

It also broke a few tests lol

[ricardoV94]

Interestingly enough we have been considered defining constants as zero input Ops to make code more clean in rewrites (we always have to worry about whether a variable has an owner, which constants and root inputs do not).

[Ch0ronomato]

Interestingly enough we have been considered defining constants as zero input Ops to make code more clean in rewrites (we always have to worry about whether a variable has an owner, which constants and root inputs do not).

That makes a lot of sense to me. I would be happy to sketch that out.

[Ch0ronomato]

this unfortunately suffers from the same problem some of the generated functions were getting where pytorch loses their reference somewhere in the process. We'll need to make sure it gets added to the warpper::gen_functors list

[Ch0ronomato]

Hopefully we don't need this and it was an over optimization. Try to refactor away the continue.