use scipy to avoid custom dask/numpy handling for softmax

Author: FynnBe

bioimageio/core/proc_ops.py

@@ -13,6 +13,7 @@
 )
 
 import numpy as np
+import scipy  # pyright: ignore[reportMissingTypeStubs]
 import xarray as xr
 from typing_extensions import Self, assert_never
 
@@ -538,11 +539,10 @@ class Softmax(_SimpleOperator):
     axis: AxisId = AxisId("channel")
 
     def _apply(self, x: Tensor, stat: Stat) -> Tensor:
-        x_max = x.data.max(dim=self.axis, keepdims=False)
-        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)
+        axis_idx = x.dims.index(self.axis)
+        result = scipy.special.softmax(x.data, axis=axis_idx)
+        result_xr = xr.DataArray(result, dims=x.dims)
+        return Tensor.from_xarray(result_xr)
 
     @property
     def required_measures(self) -> Collection[Measure]:

bioimageio/core/tensor.py

@@ -45,6 +45,7 @@
 
 
 # TODO: complete docstrings
+# TODO: in the long run---with improved typing in xarray---we should probably replace `Tensor` with xr.DataArray
 class Tensor(MagicTensorOpsMixin):
     """A wrapper around an xr.DataArray for better integration with bioimageio.spec
     and improved type annotations."""

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "pydantic-settings>=2.5,<3",
     "pydantic>=2.7.0,<3",
     "ruyaml",
+    "scipy",
     "tqdm",
     "typing-extensions",
     "xarray>=2023.01,<2025.3.0",
@@ -60,7 +61,6 @@ 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

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+import scipy  # pyright: ignore[reportMissingTypeStubs]
 import xarray as xr
 from typing_extensions import TypeGuard
 
@@ -396,9 +397,7 @@ def test_softmax(tid: MemberId):
     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]
-
+def test_softmax_with_scipy(tid: MemberId):
     from bioimageio.core.proc_ops import Softmax
 
     shape = (3, 32, 32)