Simplify nanmean with logical not

jax/numpy/lax_numpy.py

@@ -1228,6 +1228,26 @@ def nan_reduction(a, axis=None, out=None, keepdims=False, **kwargs):
 nansum = _make_nan_reduction(onp.nansum, sum, 0, nan_if_all_nan=False)
 nanprod = _make_nan_reduction(onp.nanprod, prod, 1, nan_if_all_nan=False)
 
+<<<<<<< HEAD
+=======
+@_wraps(onp.nanmean)
+def nanmean(a, axis=None, dtype=None, out=None, keepdims=False):
+  if out is not None:
+    raise ValueError("nanmean does not support the `out` argument.")
+  # Check and Count all non-NaN values
+  nan_mask = logical_not(isnan(a))
+  normalizer = sum(nan_mask, axis=axis, dtype=int32, keepdims=keepdims)
+  normalizer = lax.convert_element_type(normalizer, dtype)
+  #Perform mean calculation
+  if dtype is None:
+    if (onp.issubdtype(lax._dtype(a), onp.bool_) or
+        onp.issubdtype(lax._dtype(a), onp.integer)):
+      dtype = xla_bridge.canonicalize_dtype(onp.float64)
+    else:
+      dtype = lax._dtype(a)
+  td = true_divide(nansum(a, axis, dtype=dtype, keepdims=keepdims), normalizer)
+  return lax.convert_element_type(td, dtype)
+>>>>>>> 7c6f468... Simplify nanmean with logical not
 
 def _make_cumulative_reduction(onp_reduction, window_reduce, init_val,
                                squash_nan=False):