Add for loop based scalar loop

Author: Ch0ronomato

pytensor/link/pytorch/dispatch/scalar.py

@@ -5,6 +5,7 @@
     Cast,
     ScalarOp,
 )
+from pytensor.scalar.loop import ScalarLoop
 
 
 @pytorch_funcify.register(ScalarOp)
@@ -49,3 +50,23 @@ def cast(x):
         return x.to(dtype=dtype)
 
     return cast
+
+@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, eye
 from pytensor.tensor.type import matrices, matrix, scalar, vector
 
@@ -374,3 +376,18 @@ def inner_fn(x):
     f = function([x], out, mode="PYTORCH")
     f(torch.ones(3))
     assert "inner_fn" not in dir(m), "function call reference leaked"
+
+
+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)