xarray 2023 compatibility

Author: FynnBe

bioimageio/core/_resource_tests.py

@@ -133,7 +133,7 @@ def enable_determinism(
         try:
             os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
             try:
-                import tensorflow as tf
+                import tensorflow as tf  # pyright: ignore[reportMissingTypeStubs]
             except ImportError:
                 pass
             else:
@@ -147,7 +147,7 @@ def enable_determinism(
     if weight_formats is None or "keras_hdf5" in weight_formats:
         try:
             try:
-                import keras
+                import keras  # pyright: ignore[reportMissingTypeStubs]
             except ImportError:
                 pass
             else:

bioimageio/core/backends/keras_backend.py

@@ -5,7 +5,10 @@
 from numpy.typing import NDArray
 
 from bioimageio.spec._internal.io import download
-from bioimageio.spec._internal.type_guards import is_list, is_tuple
+from bioimageio.spec._internal.type_guards import (  # pyright: ignore[reportPrivateImportUsage]
+    is_list,
+    is_tuple,
+)
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.model.v0_5 import Version
 
@@ -18,14 +21,14 @@
 # by default, we use the keras integrated with tensorflow
 # TODO: check if we should prefer keras
 try:
-    import tensorflow as tf
-    from tensorflow import (
+    import tensorflow as tf  # pyright: ignore[reportMissingTypeStubs]
+    from tensorflow import (  # pyright: ignore[reportMissingTypeStubs]
         keras,  # pyright: ignore[reportUnknownVariableType,reportAttributeAccessIssue]
     )
 
     tf_version = Version(tf.__version__)
 except Exception:
-    import keras
+    import keras  # pyright: ignore[reportMissingTypeStubs]
 
     tf_version = None
 

bioimageio/core/backends/onnx_backend.py

@@ -2,10 +2,13 @@
 import warnings
 from typing import Any, List, Optional, Sequence, Union
 
-import onnxruntime as rt
+import onnxruntime as rt  # pyright: ignore[reportMissingTypeStubs]
 from numpy.typing import NDArray
 
-from bioimageio.spec._internal.type_guards import is_list, is_tuple
+from bioimageio.spec._internal.type_guards import (  # pyright: ignore[reportPrivateImportUsage]
+    is_list,
+    is_tuple,
+)
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.utils import download
 

bioimageio/core/backends/pytorch_backend.py

@@ -11,7 +11,11 @@
 from torch import nn
 from typing_extensions import assert_never
 
-from bioimageio.spec._internal.type_guards import is_list, is_ndarray, is_tuple
+from bioimageio.spec._internal.type_guards import (  # pyright: ignore[reportPrivateImportUsage]
+    is_list,
+    is_ndarray,
+    is_tuple,
+)
 from bioimageio.spec.common import ZipPath
 from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5
 from bioimageio.spec.utils import download

bioimageio/core/backends/tensorflow_backend.py

@@ -2,7 +2,7 @@
 from typing import Any, Optional, Sequence, Union
 
 import numpy as np
-import tensorflow as tf
+import tensorflow as tf  # pyright: ignore[reportMissingTypeStubs]
 from loguru import logger
 from numpy.typing import NDArray
 

bioimageio/core/backends/torchscript_backend.py

@@ -6,7 +6,10 @@
 import torch
 from numpy.typing import NDArray
 
-from bioimageio.spec._internal.type_guards import is_list, is_tuple
+from bioimageio.spec._internal.type_guards import (  # pyright: ignore[reportPrivateImportUsage]
+    is_list,
+    is_tuple,
+)
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.utils import download
 

bioimageio/core/io.py

@@ -14,7 +14,7 @@
     Union,
 )
 
-import h5py
+import h5py  # pyright: ignore[reportMissingTypeStubs]
 import numpy as np
 from imageio.v3 import imread, imwrite  # type: ignore
 from loguru import logger

bioimageio/core/stat_calculators.py

@@ -139,7 +139,13 @@ def compute(
         else:
             n = int(np.prod([tensor.sizes[d] for d in self._axes]))
 
-        var = xr.dot(c, c, dim=self._axes) / n
+        if xr.__version__.startswith("2023"):
+            var = (  # pyright: ignore[reportUnknownVariableType]
+                xr.dot(c, c, dims=self._axes) / n
+            )
+        else:
+            var = xr.dot(c, c, dim=self._axes) / n
+
         assert isinstance(var, xr.DataArray)
         std = np.sqrt(var)
         assert isinstance(std, xr.DataArray)

bioimageio/core/weight_converters/keras_to_tensorflow.py

@@ -4,7 +4,7 @@
 from typing import Union, no_type_check
 from zipfile import ZipFile
 
-import tensorflow
+import tensorflow  # pyright: ignore[reportMissingTypeStubs]
 
 from bioimageio.spec._internal.io import download
 from bioimageio.spec._internal.version_type import Version
@@ -24,7 +24,7 @@
     from tensorflow import keras  # type: ignore
 except Exception:
     # if the above fails try to export with the standalone keras
-    import keras
+    import keras  # pyright: ignore[reportMissingTypeStubs]
 
 
 def convert(

dev/env-py38.yaml

@@ -39,6 +39,6 @@ dependencies:
   - torchvision
   - tqdm
   - typing-extensions
-  - xarray>=2024.01,<2025.3.0
+  - xarray>=2023.01,<2025.3.0
   - pip:
       - -e ..