fix(jax-backend): fix failing test for mean (#27084)

ivy-llc · Nov 1, 2023 · 6e1f438 · 6e1f438
1 parent 4626cb1
commit 6e1f438
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/ivy/functional/backends/jax/statistical.py b/ivy/functional/backends/jax/statistical.py
@@ -37,7 +37,10 @@ def max(
     return jnp.max(a=jnp.asarray(x), axis=axis, keepdims=keepdims)
 
 
-@with_unsupported_dtypes({"0.4.14 and below": "bfloat16"}, backend_version)
+@with_unsupported_dtypes(
+    {"0.4.19 and below": "bfloat16"},
+    backend_version,
+)
 def mean(
     x: JaxArray,
     /,
@@ -47,7 +50,7 @@ def mean(
     out: Optional[JaxArray] = None,
 ) -> JaxArray:
     axis = tuple(axis) if isinstance(axis, list) else axis
-    return jnp.mean(x, axis=axis, keepdims=keepdims)
+    return jnp.mean(x, axis=axis, keepdims=keepdims, dtype=x.dtype)
 
 
 def _infer_dtype(dtype: jnp.dtype):

diff --git a/ivy/functional/backends/numpy/statistical.py b/ivy/functional/backends/numpy/statistical.py
@@ -54,9 +54,7 @@ def mean(
     out: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     axis = tuple(axis) if isinstance(axis, list) else axis
-    return ivy.astype(
-        np.mean(x, axis=axis, keepdims=keepdims, out=out), x.dtype, copy=False
-    )
+    return np.mean(x, axis=axis, keepdims=keepdims, dtype=x.dtype, out=out)
 
 
 mean.support_native_out = True