zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-08 01:54:37 +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,