fix typing

Author: FynnBe

src/bioimageio/core/sample.py

@@ -101,7 +101,7 @@ def batch_multi_index(self) -> Optional["pd.MultiIndex"]:
             return None
 
         for tensor in self.members.values():
-            idx = tensor.data.indexes.get(AxisId("batch"))
+            idx = tensor.data.indexes.get(AxisId("batch"))  # pyright: ignore[reportUnknownVariableType]
             if isinstance(idx, pd.MultiIndex):
                 return idx
 
@@ -316,7 +316,7 @@ def transpose(
                 (
                     batch_lengths := {
                         t.sizes[AxisId("batch")]
-                        for m, t in members.items()
+                        for t in members.values()
                         if AxisId("batch") in t.dims
                     }
                 )
@@ -335,7 +335,7 @@ def assign_batch_multi_index(self, multi_index: "pd.MultiIndex") -> Self:
         Raises:
             ValueError: If not all sample members have a batch dimension.
         """
-        if (
+        if len(
             no_batch := [
                 m for m, t in self.members.items() if AxisId("batch") not in t.dims
             ]
@@ -365,10 +365,14 @@ def unstack_batch_multi_index(
             A new `Sample` with unstacked batch multi-index for all members.
         """
         if (
-            no_batch := [
-                m for m, t in self.members.items() if AxisId("batch") not in t.dims
-            ]
-        ) == len(self.members) and errors == "raise":
+            len(
+                no_batch := [
+                    m for m, t in self.members.items() if AxisId("batch") not in t.dims
+                ]
+            )
+            == len(self.members)
+            and errors == "raise"
+        ):
             raise ValueError(f"No member has a batch dimension: {no_batch}")
 
         members = {
@@ -381,7 +385,7 @@ def unstack_batch_multi_index(
             len(
                 batch_lengths := {
                     t.sizes.get(AxisId("batch"))
-                    for m, t in members.items()
+                    for t in members.values()
                     if AxisId("batch") in t.dims
                 }
             )
@@ -391,10 +395,10 @@ def unstack_batch_multi_index(
                 f"Different batch lengths after unstacking: {batch_lengths}"
             )
 
-        stat = {
+        stat: Stat = {
             k: v.unstack_batch_multi_index(errors="ignore")
             if isinstance(v, Tensor)
-            else v
+            else float(v)
             for k, v in self.stat.items()
         }
         return self.__class__(

tests/test_tensor.py

@@ -82,7 +82,7 @@ def test_transpose_stack_extra_dims_to_batch_creates_multi_index():
     assert transposed.dims == (AxisId("batch"), AxisId("x"))
     assert transposed.sizes[AxisId("batch")] == 6
     assert transposed.sizes[AxisId("x")] == 4
-    batch_index = transposed.data.indexes[AxisId("batch")]
+    batch_index = transposed.data.indexes[AxisId("batch")]  # pyright: ignore[reportUnknownVariableType]
     assert isinstance(batch_index, pd.MultiIndex)
     assert tuple(str(level) for level in batch_index.names) == ("original_batch", "z")