Add functions and function calls to the pytensor language

[aseyboldt]

It seems to me that most reasons why anyone would clone graphs is to change the input of a part of a graph. Pretty much anywhere else we do something like this using a function call. I'm guessing that a lot of things would get easier if pytensor had support for functions, and I don't think they should be too difficult to add. Something along the lines of this might work for instance:

import pytensor
import pytensor.tensor as pt
from pytensor.graph import Type, Variable, Op, Apply, FunctionGraph
from pytensor.graph.op import HasInnerGraph
from typing import Tuple


class FunctionVariable(Variable):
    #function_type: FunctionType

    def __call__(self, *inputs):
        assert isinstance(self.owner.op, DefineFunction)
        return Call()(self, *inputs)
    
    def pushforward(self, *, return_value=True):
        """Return a function that computes the pushforward, or value and pushforward."""
        pass
    
    def pullback(self):
        """Return a function that computes the forward value
        and a second function that computes the pullback."""
        pass


class FunctionType(Type):
    inputs: Tuple[Type]
    outputs: Tuple[Type]  # If we had tuple types we could also just return Type
    variable_type = FunctionVariable

    def __init__(self, inputs, outputs):
        self.inputs = inputs
        self.outputs = outputs

    def filter(self, *args, **kwargs):
        return NotImplementedError()


# TODO add HasInnerGraph
class DefineFunction(Op):
    def __init__(self, fgraph):
        self.fgraph = fgraph

    def make_node(self, *closure_values):
        inputs = self.fgraph.inputs[len(closure_values):]
        outputs = self.fgraph.outputs
        function_type = FunctionType(
            [input.type for input in inputs],
            [output.type for output in outputs]
        )
        return Apply(self, closure_values, [function_type()])

    def perform(self, node, inputs, out):
        closure_vars = inputs
        fn = pytensor.function(self.fgraph.inputs, self.fgraph.outputs, mode="FAST_COMPILE")
        def inner(*inputs):
            output = fn(*[*closure_vars, *inputs])
            return output

        out[0][0] = inner


class Call(Op):
    def __init__(self):
        pass

    def make_node(self, fn, *inputs):
        return Apply(
            self,
            [fn, *inputs],
            [output_type() for output_type in fn.type.outputs]
        )

    def perform(self, node, inputs, out):
        fn, *inputs = inputs
        output = fn(*inputs)
        for i, output_value in enumerate(output):
            out[i][0] = output_value



def define_function(inputs, outputs, *, provide=None):
    if provide is None:
        provide = []
    fgraph = FunctionGraph([*provide, *inputs], outputs)
    return DefineFunction(fgraph)(*provide)

Usage could look something like this:

x1 = pt.dvector()
x2 = pt.dvector()

y = 2 * x2

z = x ** 2 * y

fn = define_function([x], [z], provide=[y])

v = pt.dvector()  # could also just use x

# We can apply the function multiple times with different input values for x:
out = fn(3 * v) + fn(2 * v)
func = pytensor.function([v, x2], [out], mode="FAST_COMPILE")
pytensor.dprint(func)

Elemwise{add,no_inplace} [id A] 6
 |Call [id B] 5
 | |DefineFunction [id C] 2
 | | |Elemwise{mul,no_inplace} [id D] 1
 | |   |TensorConstant{(1,) of 2.0} [id E]
 | |   |<TensorType(float64, (?,))> [id F]
 | |Elemwise{mul,no_inplace} [id G] 4
 |   |TensorConstant{(1,) of 3.0} [id H]
 |   |<TensorType(float64, (?,))> [id I]
 |Call [id J] 3
   |DefineFunction [id C] 2
   |Elemwise{mul,no_inplace} [id K] 0
     |TensorConstant{(1,) of 2.0} [id E]
     |<TensorType(float64, (?,))> [id I]

[ricardoV94]

How is that different form what OpFromGraph does? I think we should make it more of a first class citizen and perhaps give it a better name.

[ricardoV94]

Here is roughly your first example:

from functools import partial
import pytensor
import pytensor.tensor as pt
from pytensor.compile.builders import OpFromGraph

x1 = pt.dvector("x1")
x2 = pt.dvector("x2")

y = 2 * x2
y.name = "y"

z = x1 ** 2 * y
z.name = "z"

op = OpFromGraph([y, x1], [z])
op = partial(op, y)
out = op(3 * x1) + op(2 * x1)

func = pytensor.function([x1, x2], [out])
pytensor.dprint(func)

Elemwise{Add}[(0, 0)] [id A] 5
 |OpFromGraph{inline=False} [id B] 4
 | |Elemwise{mul,no_inplace} [id C] 'y' 1
 | | |TensorConstant{(1,) of 2.0} [id D]
 | | |x2 [id E]
 | |Elemwise{mul,no_inplace} [id F] 3
 |   |TensorConstant{(1,) of 3.0} [id G]
 |   |x1 [id H]
 |OpFromGraph{inline=False} [id I] 2
   |Elemwise{mul,no_inplace} [id C] 'y' 1
   |Elemwise{mul,no_inplace} [id J] 0
     |TensorConstant{(1,) of 2.0} [id D]
     |x1 [id H]

Inner graphs:
OpFromGraph{inline=False} [id B]
 >Elemwise{mul,no_inplace} [id K] 'z'
 > |Elemwise{pow,no_inplace} [id L]
 > | |*1-<TensorType(float64, (?,))> [id M]
 > | |InplaceDimShuffle{x} [id N]
 > |   |TensorConstant{2} [id O]
 > |*0-<TensorType(float64, (?,))> [id P]
OpFromGraph{inline=False} [id I]
 >Elemwise{mul,no_inplace} [id K] 'z'

Which reminds me, we could add a helper that collects missing inputs and assume they are constant, like we do with scan. And we should add some rewrites to pull common computations out of multiple OpFromGraph (or merge them), again like we do with Scan.

[lucianopaz]

I agree, it looks like an OpFromGraph, unless the function has a way to handle updates of the inner graph

[aseyboldt]

I think you are right, in that OpFromGraph is really similar to this. I don't think it is the same however. Whether it is different and useful enough to make the change I don't know yet.

I guess the main difference is that with the current OpFromGraph approach we do not get FunctionVariables.
Let's say we implement an Op that does integration. In python this might look like this:

def integrate(func, x0, x1, ...):
    pass

If we assume for a moment that an OpFromGraph instance is our idea of a function, then we can not write a "function" like that for integrate. One of the inputs is a function itself, so it would have to be one of the inputs in the inner FunctionGraph, but all inputs have to be subtypes of Variable, but func is itself a function, ie an instance of OpFromGraph, so it would have type Op.

In contrast, if functions are the output of DefineFunction, then they are Variables (FunctionVariables to be precise), and we could use those as input for our integrate function.

There are also a lot of other things that I think we could express using function variables, but not with OpFromGraph. I think we might want to disallow those however, because while I think they would be quite useful, I don't really know how easy they would be to implement in various backends. (And to disallow them I think all we'd have to do is prevent the outputs of all FunctionGraphs to contain a function variable. For instance:

• Closures:

  x = dscalar()
  z = dscalar()
  foo = define_function([x], [define_function([z], [2 * x * z], provide=[x])])
  # And call the function like
  integrate(foo(3.), 0, 1)
  # which would correspond to something like
  def foo(x):
      def bar(z):
          return 2 * x * z
      return bar
  integrate(foo(3), 0, 1)

• If statements involving functions:

  x = pt.dscalar()
  func = with pt.if_else(x > 0, x) as (if_, otherwise):
      y = pt.dscalar()
      if_(MakeFunction([y], [y ** x])  #if x > 0, return function 1
      otherwise(MakeFunction([y], [y ** 2]))  # if x <= 0 return function 2
  out = integrate(func, 0, 1)  # integrate the returned function

Both of those would only work by returning a function variable from an inner graph, and if we don't know how to handle cases like this in backends it is easy enough to check this when a FunctionGraph is created.

Having function variables would also be much closer to how eg llvm (and I think pretty much all other compilers?) handles this. I think our Op instances are pretty similar to instructions in llvm. And functions in llvm are not instructions (you don't define a new LLVM ir instruction every time you define a function), but they are variables that are called using a call instruction. Ops are the basic building blocks of the language, functions are something larger and user-defined that is made using those building blocks.

[ricardoV94]

Right. Usually in those cases (i.e. Integrate and Switch) we would expect users to provide us something with which we can build a FunctionGraph with. This could be a FunctionGraph directly (e.g., when building Composites), a list of inputs and outputs (when building an OpFromGraph or FunctionGraph) or a function which will return the outputs given dummy inputs generated by us (Scan). It makes sense to standardize a bit across all these APIs perhaps.

I am not sure we need a new object altogether, but I think we can definitely design a more consistent API