Add filter helper

Author: ricardoV94

pytensor/loop/basic.py

@@ -118,3 +118,35 @@ def reduce(
     if len(final_states) == 1:
         return final_states[0]
     return final_states
+
+
+def filter(
+    fn,
+    sequences,
+    non_sequences=None,
+    go_backwards=False,
+):
+    if not isinstance(sequences, (tuple, list)):
+        sequences = [sequences]
+
+    _, masks = scan(
+        fn=fn,
+        sequences=sequences,
+        non_sequences=non_sequences,
+        go_backwards=go_backwards,
+    )
+
+    if not all(mask.dtype == "bool" for mask in masks):
+        raise TypeError("The output of filter fn should be a boolean variable")
+    if len(masks) == 1:
+        masks = [masks[0]] * len(sequences)
+    elif len(masks) != len(sequences):
+        raise ValueError(
+            "filter fn must return one variable or len(sequences), but it returned {len(masks)}"
+        )
+
+    filtered_sequences = [seq[mask] for seq, mask in zip(sequences, masks)]
+
+    if len(filtered_sequences) == 1:
+        return filtered_sequences[0]
+    return filtered_sequences

tests/loop/basic.py

@@ -1,8 +1,8 @@
 import numpy as np
 
 import pytensor
-from pytensor.loop.basic import map, reduce, scan
-from pytensor.tensor import vector, zeros
+from pytensor.loop.basic import filter, map, reduce, scan
+from pytensor.tensor import eq, vector, zeros
 
 
 def test_scan():
@@ -38,3 +38,12 @@ def test_reduce():
     np.testing.assert_almost_equal(
         y.eval({xs: np.arange(10)}), np.arange(10).cumsum()[-1]
     )
+
+
+def test_filter():
+    xs = vector("xs")
+    ys = filter(
+        fn=lambda x: eq(x % 2, 0),
+        sequences=xs,
+    )
+    np.testing.assert_array_equal(ys.eval({xs: np.arange(0, 20)}), np.arange(0, 20, 2))