zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-11-03 09:58:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Added clip function (#159)

Browse Source 
* Added clip

* Added Python bindings

* Formatting

* Added cpp tests

* Added Python tests

* python bindings work

* rebase

---------

Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
Cyril Zakka, MD
2023-12-17 22:00:29 -06:00
committed by GitHub
parent ee0c2835c5
commit 8eb56beb3a
6 changed files with 116 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
python/src/ops.cpp
View File

@@ -2353,6 +2353,45 @@ void init_ops(py::module_& m) {
Returns:
array: The resulting stacked array.
)pbdoc");
m.def(
"clip",
[](const array& a,
const std::optional<ScalarOrArray>& min,
const std::optional<ScalarOrArray>& max,
StreamOrDevice s) {
std::optional<array> min_ = std::nullopt;
std::optional<array> max_ = std::nullopt;
if (min) {
min_ = to_array(min.value());
}
if (max) {
max_ = to_array(max.value());
}
return clip(a, min_, max_, s);
},
"a"_a,
py::pos_only(),
"a_min"_a,
"a_max"_a,
py::kw_only(),
"stream"_a = none,
R"pbdoc(
clip(a: array, /, a_min: Union[scalar, array, None], a_max: Union[scalar, array, None], *, stream: Union[None, Stream, Device] = None) -> array
Clip the values of the array between the given minimum and maximum.
If either ``a_min`` or ``a_max`` are ``None``, then corresponding edge
is ignored. At least one of ``a_min`` and ``a_max`` cannot be ``None``.
The input ``a`` and the limits must broadcast with one another.
Args:
a (array): Input array.
a_min (scalar or array or None): Minimum value to clip to.
a_max (scalar or array or None): Maximum value to clip to.
Returns:
array: The clipped array.
)pbdoc");
m.def(
"pad",
[](const array& a,