zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-27 19:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: enhancement of the error message for mlx.core.mean (#608)

Browse Source 
* add error message
This commit is contained in:
Avikant Srivastava 2024-02-05 14:51:49 +05:30 committed by GitHub
parent e319383ef9
commit 31fea3758e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/ops.cpp
View File

@ -1319,6 +1319,15 @@ array mean(
const std::vector<int>& axes,
bool keepdims /* = false */,
StreamOrDevice s /* = {}*/) {
int ndim = a.ndim();
for (int axis : axes) {
if (axis < -ndim || axis >= ndim) {
std::ostringstream msg;
msg << "[mean] axis " << axis + " is out of bounds for array with "
<< ndim + " dimensions.";
throw std::invalid_argument(msg.str());
}
}
auto nelements = compute_number_of_elements(a, axes);
auto dtype = at_least_float(a.dtype());
return multiply(sum(a, axes, keepdims, s), array(1.0 / nelements, dtype), s);