zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

add median op (#2705)

Browse Source
This commit is contained in:
Awni Hannun
2025-10-27 11:33:42 -07:00
committed by GitHub
parent c4767d110f
commit 539d8322d1
5 changed files with 164 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
mlx/ops.h
View File

@@ -539,6 +539,26 @@ array mean(
bool keepdims = false,
StreamOrDevice s = {});
/** Computes the median of the elements of an array. */
array median(const array& a, bool keepdims, StreamOrDevice s = {});
inline array median(const array& a, StreamOrDevice s = {}) {
return median(a, false, to_stream(s));
}
/** Computes the median of the elements of an array along the given axes */
array median(
const array& a,
const std::vector<int>& axes,
bool keepdims = false,
StreamOrDevice s = {});
/** Computes the median of the elements of an array along the given axis */
array median(
const array& a,
int axis,
bool keepdims = false,
StreamOrDevice s = {});
/** Computes the variance of the elements of an array. */
array var(const array& a, bool keepdims, int ddof = 0, StreamOrDevice s = {});
inline array var(const array& a, StreamOrDevice s = {}) {