zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-18 15:28:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add support for repeat (#278)

Browse Source 
* add repeat function

* fix styling

* optimizing repeat

* fixed minor issues

* not sure why that folder is there xD

* fixed now for sure

* test repeat not repeat test

* Fixed

---------

Co-authored-by: Bahaa Eddin tabbakha <bahaa@Bahaas-MacBook-Pro.local>
This commit is contained in:
Bahaa
2023-12-28 00:11:38 +03:00
committed by GitHub
parent 4417e37ede
commit ff2b58e299
6 changed files with 149 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
python/src/ops.cpp
View File

@@ -2405,6 +2405,40 @@ void init_ops(py::module_& m) {
Returns:
array: The resulting stacked array.
)pbdoc");
m.def(
"repeat",
[](const array& array,
int repeats,
std::optional<int> axis,
StreamOrDevice s) {
if (axis.has_value()) {
return repeat(array, repeats, axis.value(), s);
} else {
return repeat(array, repeats, s);
}
},
"array"_a,
py::pos_only(),
"repeats"_a,
"axis"_a = none,
py::kw_only(),
"stream"_a = none,
R"pbdoc(
repeat(array: array, repeats: int, axis: Optional[int] = None, *, stream: Union[None, Stream, Device] = None) -> array
Repeate an array along a specified axis.
Args:
array (array): Input array.
repeats (int): The number of repetitions for each element.
axis (int, optional): The axis in which to repeat the array along. If
unspecified it uses the flattened array of the input and repeates
along axis 0.
stream (Stream, optional): Stream or device. Defaults to ``None``.
Returns:
array: The resulting repeated array.
)pbdoc");
m.def(
"clip",
[](const array& a,

31
python/tests/test_ops.py
View File

@@ -1516,6 +1516,37 @@ class TestOps(mlx_tests.MLXTestCase):
expected = mx.array(np.linspace(0, 1, 10))
self.assertEqualArray(d, expected)
def test_repeat(self):
# Setup data for the tests
data = np.array([[[13, 3], [16, 6]], [[14, 4], [15, 5]], [[11, 1], [12, 2]]])
# Test repeat along axis 0
repeat_axis_0 = mx.repeat(mx.array(data), 2, axis=0)
expected_axis_0 = np.repeat(data, 2, axis=0)
self.assertEqualArray(repeat_axis_0, mx.array(expected_axis_0))
# Test repeat along axis 1
repeat_axis_1 = mx.repeat(mx.array(data), 2, axis=1)
expected_axis_1 = np.repeat(data, 2, axis=1)
self.assertEqualArray(repeat_axis_1, mx.array(expected_axis_1))
# Test repeat along the last axis (default)
repeat_axis_2 = mx.repeat(mx.array(data), 2)
expected_axis_2 = np.repeat(data, 2)
self.assertEqualArray(repeat_axis_2, mx.array(expected_axis_2))
# Test repeat with a 1D array along axis 0
data_2 = mx.array([1, 3, 2])
repeat_2 = mx.repeat(mx.array(data_2), 3, axis=0)
expected_2 = np.repeat(data_2, 3)
self.assertEqualArray(repeat_2, mx.array(expected_2))
# Test repeat with a 2D array along axis 0
data_3 = mx.array([[1, 2, 3], [4, 5, 4], [0, 1, 2]])
repeat_3 = mx.repeat(mx.array(data_3), 2, axis=0)
expected_3 = np.repeat(data_3, 2, axis=0)
self.assertEqualArray(repeat_3, mx.array(expected_3))
if __name__ == "__main__":
unittest.main()