fix: corrected implementation of where param numpy mean

ivy/functional/frontends/numpy/statistics/averages_and_variances.py

@@ -72,62 +72,30 @@ def cov(
 @handle_numpy_dtype
 @to_ivy_arrays_and_back
 @from_zero_dim_arrays_to_scalar
-def mean(
-    a,
-    /,
-    *,
-    axis=None,
-    keepdims=False,
-    out=None,
-    dtype=None,
-    where=True,
-):
+def mean(a, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
     axis = tuple(axis) if isinstance(axis, list) else axis
-    if dtype:
-        a = ivy.astype(ivy.array(a), ivy.as_ivy_dtype(dtype))
-
-    ret = ivy.mean(a, axis=axis, keepdims=keepdims, out=out)
-    if ivy.is_array(where):
-        ret = ivy.where(where, ret, ivy.default(out, ivy.zeros_like(ret)), out=out)
+    dtype = dtype or a.dtype if not ivy.is_int_dtype(a.dtype) else ivy.float64
+    where = ivy.where(where, ivy.ones_like(a), 0)
+    if where is not True:
+        a = ivy.where(where, a, 0.0)
+        sum = ivy.sum(a, axis=axis, keepdims=keepdims, dtype=dtype)
+        cnt = ivy.sum(where, axis=axis, keepdims=keepdims, dtype=int)
+        ret = ivy.divide(sum, cnt, out=out)
+    else:
+        ret = ivy.mean(a.astype(dtype), axis=axis, keepdims=keepdims, out=out)
 
-    return ret
+    return ret.astype(dtype)
 
 
 @handle_numpy_out
 @handle_numpy_dtype
 @to_ivy_arrays_and_back
 @from_zero_dim_arrays_to_scalar
-def nanmean(
-    a,
-    /,
-    *,
-    axis=None,
-    keepdims=False,
-    out=None,
-    dtype=None,
-    where=True,
-):
-    is_nan = ivy.isnan(a)
-    axis = tuple(axis) if isinstance(axis, list) else axis
-
-    if not ivy.any(is_nan):
-        if dtype:
-            a = ivy.astype(ivy.array(a), ivy.as_ivy_dtype(dtype))
-        ret = ivy.mean(a, axis=axis, keepdims=keepdims, out=out)
-
-        if ivy.is_array(where):
-            ret = ivy.where(where, ret, ivy.default(out, ivy.zeros_like(ret)), out=out)
-
-    else:
-        a = [i for i in a if ivy.isnan(i) is False]
-
-        if dtype:
-            a = ivy.astype(ivy.array(a), ivy.as_ivy_dtype(dtype))
-        ret = ivy.mean(a, axis=axis, keepdims=keepdims, out=out)
-
-        if ivy.is_array(where):
-            ret = ivy.where(where, ret, ivy.default(out, ivy.zeros_like(ret)), out=out)
-
+def nanmean(a, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
+    where = ~ivy.isnan(a) & where
+    ret = mean(a, axis, dtype, keepdims=keepdims, where=where).ivy_array
+    if out is not None:
+        out.data = ret.data
     return ret
 
 

ivy_tests/test_ivy/test_frontends/test_numpy/test_statistics/test_averages_and_variances.py

@@ -223,32 +223,28 @@ def test_numpy_mean(
     keep_dims,
 ):
     input_dtypes, x, axis = dtype_and_x
-    if isinstance(axis, tuple):
-        axis = axis[0]
-
     where, input_dtypes, test_flags = np_frontend_helpers.handle_where_and_array_bools(
         where=where,
         input_dtype=input_dtypes,
         test_flags=test_flags,
     )
-
-    np_frontend_helpers.test_frontend_function(
-        input_dtypes=input_dtypes,
-        frontend=frontend,
-        backend_to_test=backend_fw,
-        test_flags=test_flags,
-        fn_tree=fn_tree,
-        on_device=on_device,
-        atol=1e-2,
-        rtol=1e-2,
-        x=x[0],
-        axis=axis,
-        dtype=dtype[0],
-        out=None,
-        keepdims=keep_dims,
-        where=where,
-        test_values=False,
-    )
+    if "bfloat16" not in input_dtypes + dtype or []:
+        helpers.test_frontend_function(
+            input_dtypes=input_dtypes,
+            frontend=frontend,
+            backend_to_test=backend_fw,
+            test_flags=test_flags,
+            fn_tree=fn_tree,
+            on_device=on_device,
+            atol=1e-2,
+            rtol=1e-2,
+            a=x[0],
+            axis=axis,
+            dtype=dtype[0],
+            out=None,
+            keepdims=keep_dims,
+            where=where,
+        )
 
 
 # nanmean
@@ -279,24 +275,23 @@ def test_numpy_nanmean(
         input_dtype=input_dtypes,
         test_flags=test_flags,
     )
-
-    np_frontend_helpers.test_frontend_function(
-        input_dtypes=input_dtypes,
-        frontend=frontend,
-        backend_to_test=backend_fw,
-        test_flags=test_flags,
-        fn_tree=fn_tree,
-        on_device=on_device,
-        rtol=1e-2,
-        atol=1e-2,
-        a=a[0],
-        axis=axis,
-        dtype=dtype[0],
-        out=None,
-        keepdims=keep_dims,
-        where=where,
-        test_values=False,
-    )
+    if "bfloat16" not in input_dtypes + dtype or []:
+        helpers.test_frontend_function(
+            input_dtypes=input_dtypes,
+            frontend=frontend,
+            backend_to_test=backend_fw,
+            test_flags=test_flags,
+            fn_tree=fn_tree,
+            on_device=on_device,
+            rtol=1e-2,
+            atol=1e-2,
+            a=a[0],
+            axis=axis,
+            dtype=dtype[0],
+            out=None,
+            keepdims=keep_dims,
+            where=where,
+        )
 
 
 @handle_frontend_test(