make softmax compatible with xarray <= 2024.07

Author: FynnBe

bioimageio/core/proc_ops.py

@@ -539,7 +539,8 @@ class Softmax(_SimpleOperator):
 
     def _apply(self, x: Tensor, stat: Stat) -> Tensor:
         x_max = x.data.max(dim=self.axis, keepdims=False)
-        exp_x_shifted: xr.DataArray = xr.ufuncs.exp(x.data - x_max)
+        x_shifted = x.data - x_max
+        exp_x_shifted = xr.DataArray(x_shifted.data.exp(), dims=x.dims)
         result = exp_x_shifted / exp_x_shifted.sum(dim=self.axis)
         return Tensor.from_xarray(result)
 

pyproject.toml

@@ -60,6 +60,7 @@ dev = [
     "pytest-cov",
     "pytest",
     "python-dotenv",
+    "scipy",
     "segment-anything",  # for model testing
     "timm",              # for model testing
     "torch>=1.6,<3",

tests/test_proc_ops.py

@@ -394,3 +394,23 @@ def test_softmax(tid: MemberId):
         dims=axes,
     )
     xr.testing.assert_allclose(exp, sample.members[tid].data, rtol=1e-5, atol=1e-7)
+
+
+def test_softmax_w_scipy(tid: MemberId):
+    import scipy  # pyright: ignore[reportMissingTypeStubs]
+
+    from bioimageio.core.proc_ops import Softmax
+
+    shape = (3, 32, 32)
+    axes = ("channel", "y", "x")
+    np_data = np.random.rand(*shape)
+    data = xr.DataArray(np_data, dims=axes)
+    sample = Sample(members={tid: Tensor.from_xarray(data)}, stat={}, id=None)
+    softmax = Softmax(tid, tid, axis=AxisId("channel"))
+    softmax(sample)
+
+    exp = xr.DataArray(
+        scipy.special.softmax(np_data, axis=0),
+        dims=axes,
+    )
+    xr.testing.assert_allclose(exp, sample.members[tid].data, rtol=1e-5, atol=1e-7)