Basic support for makeop

Author: Ian Schweer

pytensor/compile/__init__.py

@@ -30,6 +30,7 @@
     OPT_O3,
     OPT_STABILIZE,
     OPT_UNSAFE,
+    PYTORCH,
     AddDestroyHandler,
     AddFeatureOptimizer,
     Mode,

pytensor/link/pytorch/dispatch/basic.py

@@ -3,6 +3,8 @@
 
 import torch
 
+from pytensor.compile import PYTORCH
+from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.ops import DeepCopyOp
 from pytensor.graph.fg import FunctionGraph
 from pytensor.link.utils import fgraph_to_python
@@ -132,3 +134,17 @@ def makevector(*x):
         return torch.tensor(x, dtype=torch_dtype)
 
     return makevector
+
+
+@pytorch_funcify.register(OpFromGraph)
+def pytorch_funcify_OpFromGraph(op, node=None, **kwargs):
+    _ = kwargs.pop("storage_map", None)
+
+    PYTORCH.optimizer(op.fgraph)
+    fgraph_fn = torch.compile(pytorch_funcify(op.fgraph, **kwargs))
+
+    def opfromgraph(*inputs, dim=op.fgraph.outputs):
+        res = fgraph_fn(*inputs)
+        return res[0]
+
+    return opfromgraph

tests/link/pytorch/test_basic.py

@@ -5,6 +5,7 @@
 import pytest
 
 import pytensor.tensor.basic as ptb
+from pytensor.compile.builders import OpFromGraph
 from pytensor.compile.function import function
 from pytensor.compile.mode import get_mode
 from pytensor.compile.sharedvalue import SharedVariable, shared
@@ -14,7 +15,7 @@
 from pytensor.graph.op import Op
 from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor import alloc, arange, as_tensor, empty, eye
-from pytensor.tensor.type import matrix, scalar, vector
+from pytensor.tensor.type import matrices, matrix, scalar, vector
 
 
 torch = pytest.importorskip("torch")
@@ -301,3 +302,25 @@ def test_pytorch_MakeVector():
     x_fg = FunctionGraph([], [x])
 
     compare_pytorch_and_py(x_fg, [])
+
+
+def test_pytorch_OpFromGraph():
+    x, y, z = matrices("xyz")
+    ofg_1 = OpFromGraph([x, y], [x + y])
+    OpFromGraph([x, y], [x * y, x - y])
+
+    # o1, o2 = ofg_2(y, z)
+    # out = ofg_1(x, o1) + o2
+
+    out = ofg_1(y, z)
+
+    xv = np.ones((2, 2), dtype=config.floatX)
+    np.ones((2, 2), dtype=config.floatX) * 3
+    zv = np.ones((2, 2), dtype=config.floatX) * 5
+
+    f = FunctionGraph([y, z], [out])
+    import pytensor.printing
+
+    pytensor.printing.debugprint(f)
+
+    compare_pytorch_and_py(f, [xv, zv])