Add test for arbitrary padding at higher dimensions

Author: jessegrabowski

tests/tensor/test_pad.py

@@ -26,7 +26,7 @@ def test_constant_pad(
     x = np.random.normal(size=size).astype(floatX)
     expected = np.pad(x, pad_width, mode="constant", constant_values=constant)
     z = pad(x, pad_width, mode="constant", constant_values=constant)
-    assert z.pad_mode == "constant"
+    assert z.owner.op.pad_mode == "constant"
 
     f = pytensor.function([], z, mode="FAST_COMPILE")
 
@@ -43,7 +43,7 @@ def test_edge_pad(size: tuple, pad_width: int | tuple[int, ...]):
     x = np.random.normal(size=size).astype(floatX)
     expected = np.pad(x, pad_width, mode="edge")
     z = pad(x, pad_width, mode="edge")
-    assert z.pad_mode == "edge"
+    assert z.owner.op.pad_mode == "edge"
 
     f = pytensor.function([], z, mode="FAST_COMPILE")
 
@@ -65,7 +65,7 @@ def test_linear_ramp_pad(
     x = np.random.normal(size=size).astype(floatX)
     expected = np.pad(x, pad_width, mode="linear_ramp", end_values=end_values)
     z = pad(x, pad_width, mode="linear_ramp", end_values=end_values)
-    assert z.pad_mode == "linear_ramp"
+    assert z.owner.op.pad_mode == "linear_ramp"
 
     f = pytensor.function([], z, mode="FAST_COMPILE")
 
@@ -89,8 +89,7 @@ def test_stat_pad(
     x = np.random.normal(size=size).astype(floatX)
     expected = np.pad(x, pad_width, mode=stat, stat_length=stat_length)
     z = pad(x, pad_width, mode=stat, stat_length=stat_length)
-    assert z.pad_mode == stat
-    assert z.stat_length_input == (stat_length is not None)
+    assert z.owner.op.pad_mode == stat
 
     f = pytensor.function([], z, mode="FAST_COMPILE")
 
@@ -107,7 +106,7 @@ def test_wrap_pad(size: tuple, pad_width: int | tuple[int, ...]):
     x = np.random.normal(size=size).astype(floatX)
     expected = np.pad(x, pad_width, mode="wrap")
     z = pad(x, pad_width, mode="wrap")
-    assert z.pad_mode == "wrap"
+    assert z.owner.op.pad_mode == "wrap"
     f = pytensor.function([], z, mode="FAST_COMPILE")
 
     np.testing.assert_allclose(expected, f(), atol=ATOL, rtol=RTOL)
@@ -128,7 +127,50 @@ def test_symmetric_pad(size, pad_width, reflect_type):
     x = np.random.normal(size=size).astype(floatX)
     expected = np.pad(x, pad_width, mode="symmetric", reflect_type=reflect_type)
     z = pad(x, pad_width, mode="symmetric", reflect_type=reflect_type)
-    assert z.pad_mode == "symmetric"
+    assert z.owner.op.pad_mode == "symmetric"
+    f = pytensor.function([], z, mode="FAST_COMPILE")
+
+    np.testing.assert_allclose(expected, f(), atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        "constant",
+        "edge",
+        "linear_ramp",
+        "wrap",
+        "symmetric",
+        "mean",
+        "maximum",
+        "minimum",
+    ],
+)
+@pytest.mark.parametrize("padding", ["symmetric", "asymmetric"])
+def test_nd_padding(mode, padding):
+    rng = np.random.default_rng()
+    n = rng.integers(3, 10)
+    if padding == "symmetric":
+        pad_width = [(i, i) for i in rng.integers(1, 5, size=n)]
+        stat_length = [(i, i) for i in rng.integers(1, 5, size=n)]
+    else:
+        pad_width = rng.integers(1, 5, size=(n, 2)).tolist()
+        stat_length = rng.integers(1, 5, size=(n, 2)).tolist()
+
+    test_kwargs = {
+        "constant": {"constant_values": 0},
+        "linear_ramp": {"end_values": 0},
+        "maximum": {"stat_length": stat_length},
+        "mean": {"stat_length": stat_length},
+        "minimum": {"stat_length": stat_length},
+        "reflect": {"reflect_type": "even"},
+        "symmetric": {"reflect_type": "even"},
+    }
+
+    x = np.random.normal(size=(2,) * n).astype(floatX)
+    kwargs = test_kwargs.get(mode, {})
+    expected = np.pad(x, pad_width, mode=mode, **kwargs)
+    z = pad(x, pad_width, mode=mode, **kwargs)
     f = pytensor.function([], z, mode="FAST_COMPILE")
 
     np.testing.assert_allclose(expected, f(), atol=ATOL, rtol=RTOL)