Add for loop based scalar loop

Author: Ch0ronomato

pytensor/link/pytorch/dispatch/scalar.py

@@ -8,6 +8,7 @@
     ScalarOp,
 )
 from pytensor.scalar.math import Softplus
+from pytensor.scalar.loop import ScalarLoop
 
 
 @pytorch_funcify.register(ScalarOp)
@@ -58,7 +59,26 @@ def cast(x):
 
     return cast
 
-
 @pytorch_funcify.register(Softplus)
 def pytorch_funcify_Softplus(op, node, **kwargs):
     return torch.nn.Softplus()
+
+@pytorch_funcify.register(ScalarLoop)
+def pytorch_funicify_ScalarLoop(op, node, **kwargs):
+    update = pytorch_funcify(op.fgraph)
+
+    def inner(steps, start, constant, update=update, is_while=op.is_while):
+        # easiest way to do it is to loop
+        c = start
+        for i in range(steps):
+            outs = update(c, constant)
+            if is_while:
+                n, done = outs
+                if done:
+                    return n
+                c = n
+            else:
+                c = outs[0]
+        return c
+
+    return inner

tests/link/pytorch/test_basic.py

@@ -17,6 +17,8 @@
 from pytensor.ifelse import ifelse
 from pytensor.link.pytorch.linker import PytorchLinker
 from pytensor.raise_op import CheckAndRaise
+from pytensor.scalar import float64, int64
+from pytensor.scalar.loop import ScalarLoop
 from pytensor.tensor import alloc, arange, as_tensor, empty, expit, eye, softplus
 from pytensor.tensor.type import matrices, matrix, scalar, vector
 
@@ -388,3 +390,17 @@ def test_pytorch_softplus():
     out = softplus(x)
     f = FunctionGraph([x], [out])
     compare_pytorch_and_py(f, [np.random.rand(3)])
+
+def test_ScalarLoop():
+    n_steps = int64("n_steps")
+    x0 = float64("x0")
+    const = float64("const")
+    x = x0 + const
+
+    op = ScalarLoop(init=[x0], constant=[const], update=[x])
+    x = op(n_steps, x0, const)
+
+    fn = function([n_steps, x0, const], x, mode=pytorch_mode)
+    np.testing.assert_allclose(fn(5, 0, 1), 5)
+    np.testing.assert_allclose(fn(5, 0, 2), 10)
+    np.testing.assert_allclose(fn(4, 3, -1), -1)