Changes to allow nan functions to work with xarray

cubed-dev · Jan 15, 2025 · 039ec69 · 039ec69
1 parent 5f75ba2
commit 039ec69
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/cubed/array_api/array_object.py b/cubed/array_api/array_object.py
@@ -367,9 +367,9 @@ def __array_namespace__(self, /, *, api_version=None):
             "2023.12",
         ):
             raise ValueError(f"Unrecognized array API version: {api_version!r}")
-        import cubed.array_api as array_api
+        import cubed
 
-        return array_api
+        return cubed
 
     def __bool__(self, /):
         if self.ndim != 0:

diff --git a/cubed/nan_functions.py b/cubed/nan_functions.py
@@ -18,9 +18,9 @@
 # https://github.com/data-apis/array-api/issues/621
 
 
-def nanmean(x, /, *, axis=None, keepdims=False, split_every=None):
+def nanmean(x, /, *, axis=None, dtype=None, keepdims=False, split_every=None):
     """Compute the arithmetic mean along the specified axis, ignoring NaNs."""
-    dtype = x.dtype
+    dtype = dtype or x.dtype
     intermediate_dtype = [("n", nxp.int64), ("total", nxp.float64)]
     return reduction(
         x,