diff --git a/autoray/autoray.py b/autoray/autoray.py index 386177e..82e5064 100644 --- a/autoray/autoray.py +++ b/autoray/autoray.py @@ -4124,7 +4124,8 @@ def torch_pad(array, pad_width, mode="constant", constant_values=0): try: # numpy takes pads like ((0, 0), (1, 1), ... (n-1, n-1)) # torch takes pads like (n-1, n-1, n-2, n-2, n-3, n-3, ...) - pad = tuple(itertools.chain.from_iterable(pad_width))[::-1] + # i.e. axes reversed, but each (before, after) pair kept in order + pad = tuple(itertools.chain.from_iterable(reversed(pad_width))) # a single tuple was specified ((a, b),) - use for all axes if len(pad) == 2: diff --git a/tests/test_autoray.py b/tests/test_autoray.py index 6614fe3..b8f88d7 100644 --- a/tests/test_autoray.py +++ b/tests/test_autoray.py @@ -418,12 +418,18 @@ def test_pad(backend): (((1, 2),), (6, 7, 8)), # different pad for every axis (((4, 3), (2, 4), (3, 2)), (10, 10, 10)), + # asymmetric, including zero pads + (((0, 1), (2, 0), (3, 1)), (4, 6, 9)), ]: B = ar.do("pad", A, pad_width) assert shape(B) == new_shape assert ar.to_numpy(ar.do("sum", A)) == pytest.approx( ar.to_numpy(ar.do("sum", B)) ) + # check the padding is placed the same way numpy places it + np.testing.assert_allclose( + ar.to_numpy(B), np.pad(ar.to_numpy(A), pad_width) + ) @pytest.mark.parametrize("backend", gen_params(backends=...))