diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md index 186908f09c..d83832be8f 100644 --- a/ACKNOWLEDGMENTS.md +++ b/ACKNOWLEDGMENTS.md @@ -20,6 +20,7 @@ MLX was developed with contributions from the following individuals: - Paul Paczuski: Improved stability of BCE loss calculation - Max-Heinrich Laves: Added `conv_transpose1d`, `conv_transpose2d`, and `conv_transpose3d` ops. - Gökdeniz Gülmez: Added the `Muon (MomentUm Orthogonalized by Newton-schulz)` optimizer, and the `ReLU²` activation function. +- katlun-lgtm: Added `reflect` and `symmetric` padding modes. diff --git a/mlx/ops.cpp b/mlx/ops.cpp index afc9ab489d..7cc283e399 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -1460,6 +1460,85 @@ array tile( return reshape(x, std::move(final_shape), s); } +array reflect_pad( + const array& a, + const std::vector& axes, + const Shape& low_pad_size, + const Shape& high_pad_size, + const Shape& out_shape, + bool include_edge, + StreamOrDevice s /* = {} */) { + array out = zeros(out_shape, a.dtype(), s); + Shape starts(a.ndim(), 0); + auto stops = a.shape(); + for (size_t i = 0; i < axes.size(); i++) { + int ax = axes[i]; + starts[ax] = low_pad_size[i]; + stops[ax] += low_pad_size[i]; + } + // Copy over values from the unpadded array + array padded = slice_update(out, a, starts, stops, s); + + for (size_t i = 0; i < axes.size(); i++) { + int ax = axes[i]; + int n = a.shape(ax); + int L = low_pad_size[i]; + int H = high_pad_size[i]; + if (L == 0 && H == 0) { + continue; + } + // reflect skips the edge value (period 2(n-1)); symmetric repeats it + // (period 2n). + int offset = (!include_edge && n > 1) ? 1 : 0; + int tile = n - offset; + + if (L > 0) { + int filled_start = low_pad_size[i]; + int remaining = L; + while (remaining > 0) { + int chunk = std::min(remaining, tile); + Shape src_starts(a.ndim(), 0); + Shape src_stops = out_shape; + src_starts[ax] = filled_start + offset; + src_stops[ax] = filled_start + offset + chunk; + array piece = flip(slice(padded, src_starts, src_stops, s), ax, s); + + Shape dst_starts(a.ndim(), 0); + Shape dst_stops = out_shape; + dst_starts[ax] = filled_start - chunk; + dst_stops[ax] = filled_start; + padded = slice_update(padded, piece, dst_starts, dst_stops, s); + + filled_start -= chunk; + remaining -= chunk; + } + } + + if (H > 0) { + int filled_end = low_pad_size[i] + n; + int remaining = H; + while (remaining > 0) { + int chunk = std::min(remaining, tile); + Shape src_starts(a.ndim(), 0); + Shape src_stops = out_shape; + src_starts[ax] = filled_end - offset - chunk; + src_stops[ax] = filled_end - offset; + array piece = flip(slice(padded, src_starts, src_stops, s), ax, s); + + Shape dst_starts(a.ndim(), 0); + Shape dst_stops = out_shape; + dst_starts[ax] = filled_end; + dst_stops[ax] = filled_end + chunk; + padded = slice_update(padded, piece, dst_starts, dst_stops, s); + + filled_end += chunk; + remaining -= chunk; + } + } + } + return padded; +} + array edge_pad( const array& a, const std::vector& axes, @@ -1552,6 +1631,12 @@ array pad( {a, astype(pad_value, a.dtype(), s)}); } else if (mode == "edge") { return edge_pad(a, axes, low_pad_size, high_pad_size, out_shape, s); + } else if (mode == "reflect") { + return reflect_pad( + a, axes, low_pad_size, high_pad_size, out_shape, false, s); + } else if (mode == "symmetric") { + return reflect_pad( + a, axes, low_pad_size, high_pad_size, out_shape, true, s); } else { std::ostringstream msg; msg << "Invalid padding mode (" << mode << ") passed to pad"; diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 0941793949..d959a69c97 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -3509,7 +3509,7 @@ void init_ops(nb::module_& m) { nb::kw_only(), "stream"_a = nb::none(), nb::sig( - "def pad(a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal['constant', 'edge'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"), + "def pad(a: array, pad_width: int | tuple[int] | tuple[int, int] | list[tuple[int, int]], mode: Literal['constant', 'edge', 'reflect', 'symmetric'] = 'constant', constant_values: scalar | array = 0, *, stream: StreamOrDevice = None) -> array"), R"pbdoc( Pad an array with a constant value @@ -3524,6 +3524,8 @@ void init_ops(nb::module_& m) { mode: Padding mode. One of the following strings: "constant" (default): Pads with a constant value. "edge": Pads with the edge values of array. + "reflect": Pads with the reflection of the array, without repeating the edge values. + "symmetric": Pads with the reflection of the array, repeating the edge values. constant_values (array or scalar, optional): Optional constant value to pad the edges of the array with. diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a0f4fadaa3..c71458bae4 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -2265,6 +2265,39 @@ def test_nan_to_num(self): out_mx = mx.nan_to_num(a, nan=0.0, posinf=1000, neginf=-1000) self.assertTrue(np.allclose(out_mx, out_np)) + def test_pad_reflect_symmetric(self): + # mx.pad reflect/symmetric must match numpy.pad exactly. Covers + # in-bounds, multi-reflect (pad larger than the axis, exercising the + # tiling loop), asymmetric + # per-axis widths, zero-width sides, and degenerate axes (n == 1, n == 2). + cases = [ + ((8,), [(2, 3)]), + ((8,), [(0, 4)]), + ((8,), [(3, 0)]), + ((8,), [(7, 8)]), + ((4,), [(10, 7)]), # multi-reflect + ((4,), [(20, 20)]), # multi-reflect, both sides + ((3,), [(9, 1)]), # multi-reflect + ((1,), [(3, 2)]), # degenerate axis + ((2,), [(5, 6)]), # smallest non-trivial, multi-reflect + ((5, 6), [(2, 3), (1, 2)]), + ((5, 6), [(9, 9), (11, 0)]), # both axes multi-reflect + ((3, 4, 5), [(1, 1), (0, 0), (2, 2)]), + ((3, 4, 5), [(4, 4), (0, 0), (7, 3)]), + ] + for mode in ("reflect", "symmetric"): + for shape, pw in cases: + a_npy = np.random.randn(*shape).astype(np.float32) + a_mlx = mx.array(a_npy) + b_npy = np.pad(a_npy, pw, mode=mode) + b_mlx = mx.pad(a_mlx, pw, mode=mode) + self.assertEqual(b_mlx.shape, tuple(b_npy.shape)) + self.assertTrue( + np.array_equal(np.array(b_mlx), b_npy), + msg=f"mismatch mode={mode} shape={shape} pad={pw}", + ) + self.assertEqual(b_mlx.dtype, mx.float32) + def test_as_strided(self): x_npy = np.random.randn(128).astype(np.float32) x_mlx = mx.array(x_npy) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 0dd9385e14..f7a2b8ab92 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -2978,6 +2978,49 @@ TEST_CASE("test pad") { 0.0f}, {4, 4}); CHECK(array_equal(padded_x, expected).item()); + + // reflect padding (mirror without repeating the edge value) + x = array({1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, {5}); + CHECK(array_equal( + pad(x, {{2, 2}}, array(0.0f), "reflect"), + array({3.0f, 2.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 4.0f, 3.0f}, {9})) + .item()); + CHECK(array_equal( + pad(x, {{0, 3}}, array(0.0f), "reflect"), + array({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 4.0f, 3.0f, 2.0f}, {8})) + .item()); + + // symmetric padding (mirror repeating the edge value) + CHECK(array_equal( + pad(x, {{2, 2}}, array(0.0f), "symmetric"), + array({2.0f, 1.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 5.0f, 4.0f}, {9})) + .item()); + CHECK(array_equal( + pad(x, {{3, 0}}, array(0.0f), "symmetric"), + array({3.0f, 2.0f, 1.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f}, {8})) + .item()); + + // multi-reflect: pad larger than the axis repeats the reflection (numpy + // parity) + x = array({1.0f, 2.0f, 3.0f}, {3}); + CHECK(array_equal( + pad(x, {{5, 5}}, array(0.0f), "reflect"), + array( + {2.0f, + 1.0f, + 2.0f, + 3.0f, + 2.0f, + 1.0f, + 2.0f, + 3.0f, + 2.0f, + 1.0f, + 2.0f, + 3.0f, + 2.0f}, + {13})) + .item()); } TEST_CASE("test power") {