Add IfElse

Author: Ian Schweer

pytensor/link/pytorch/dispatch/basic.py

@@ -4,6 +4,7 @@
 
 from pytensor.compile.ops import DeepCopyOp
 from pytensor.graph.fg import FunctionGraph
+from pytensor.ifelse import IfElse
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, Eye, Join
@@ -116,3 +117,14 @@ def eye(N, M, k):
         return zeros
 
     return eye
+
+
+@pytorch_funcify.register(IfElse)
+def pytorch_funcify_IfElse(op, **kwargs):
+    n_outs = op.n_outs
+    assert n_outs == 1
+
+    def ifelse(cond, *args, n_outs=n_outs):
+        return torch.where(cond, *args)
+
+    return ifelse

tests/link/pytorch/test_basic.py

@@ -11,7 +11,8 @@
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Apply
 from pytensor.graph.fg import FunctionGraph
-from pytensor.graph.op import Op
+from pytensor.graph.op import Op, get_test_value
+from pytensor.ifelse import ifelse
 from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor import alloc, arange, as_tensor, empty, eye
 from pytensor.tensor.type import matrix, scalar, vector
@@ -294,3 +295,20 @@ def test_eye(dtype):
         for _M in range(1, 6):
             for _k in list(range(_M + 2)) + [-x for x in range(1, _N + 2)]:
                 np.testing.assert_array_equal(fn(_N, _M, _k), np.eye(_N, _M, _k))
+
+
+def test_pytorch_ifelse():
+    true_vals = np.r_[1, 2, 3]
+    false_vals = np.r_[-1, -2, -3]
+
+    x = ifelse(np.array(True), true_vals, false_vals)
+    x_fg = FunctionGraph([], [x])
+
+    compare_pytorch_and_py(x_fg, [])
+
+    a = scalar("a")
+    a.tag.test_value = np.array(0.2, dtype=config.floatX)
+    x = ifelse(a < 0.5, true_vals, false_vals)
+    x_fg = FunctionGraph([a], [x])  # I.e. False
+
+    compare_pytorch_and_py(x_fg, [get_test_value(i) for i in x_fg.inputs])