From 31fea3758ead101e5dfdb4bd5c1d62f43598ccea Mon Sep 17 00:00:00 2001
From: Avikant Srivastava <asdavikant@gmail.com>
Date: Mon, 5 Feb 2024 14:51:49 +0530
Subject: [PATCH] feat: enhancement of the error message for mlx.core.mean
 (#608)

* add error message
---
 mlx/ops.cpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index 3f784dd9d..019624a70 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -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);