Add test for flip

Author: jessegrabowski

tests/tensor/test_pad.py

@@ -2,7 +2,7 @@
 import pytest
 
 import pytensor
-from pytensor.tensor.pad import PadMode, pad
+from pytensor.tensor.pad import PadMode, flip, pad
 
 
 floatX = pytensor.config.floatX
@@ -132,3 +132,27 @@ def test_symmetric_pad(size, pad_width, reflect_type):
     f = pytensor.function([], z, mode="FAST_COMPILE")
 
     np.testing.assert_allclose(expected, f(), atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize(
+    "size", [(3,), (3, 3), (3, 5, 5)], ids=["1d", "2d square", "3d square"]
+)
+def test_flip(size: tuple[int]):
+    from itertools import combinations
+
+    x = np.random.normal(size=size).astype(floatX)
+    x_pt = pytensor.tensor.tensor(shape=size, name="x")
+    expected = np.flip(x, axis=None)
+    z = flip(x_pt, axis=None)
+    f = pytensor.function([x_pt], z, mode="FAST_COMPILE")
+    np.testing.assert_allclose(expected, f(x), atol=ATOL, rtol=RTOL)
+
+    # Test all combinations of axes
+    flip_options = [
+        axes for i in range(1, x.ndim + 1) for axes in combinations(range(x.ndim), r=i)
+    ]
+    for axes in flip_options:
+        expected = np.flip(x, axis=list(axes))
+        z = flip(x_pt, axis=list(axes))
+        f = pytensor.function([x_pt], z, mode="FAST_COMPILE")
+        np.testing.assert_allclose(expected, f(x), atol=ATOL, rtol=RTOL)