Fix array protocol and some types

Author: mdbenito

src/pydvl/utils/array.py

@@ -73,7 +73,7 @@
     "try_torch_import",
 ]
 
-DT = TypeVar("DT")
+DT = TypeVar("DT", bound=np.generic)
 
 
 def try_torch_import(require: bool = False) -> ModuleType | None:
@@ -138,7 +138,8 @@ class Array(Protocol[DT]):
         in NumPy and PyTorch arrays.
     """
 
-    nbytes: int
+    @property
+    def nbytes(self) -> int: ...
 
     @property
     def shape(self) -> tuple[int, ...]: ...
@@ -163,7 +164,7 @@ def __mul__(self, other) -> Array: ...
 
     def __matmul__(self, other) -> Array: ...
 
-    def __array__(self, dtype: DT | None = None) -> NDArray: ...
+    def __array__(self, dtype: DT | None = None) -> NDArray[DT]: ...
 
     def flatten(self, *args, **kwargs) -> Self: ...
 

src/pydvl/valuation/samplers/classwise.py

@@ -51,7 +51,7 @@
 import numpy as np
 from more_itertools import chunked, flatten
 
-from pydvl.utils.array import Array, array_unique, is_categorical
+from pydvl.utils.array import DT, Array, array_unique, is_categorical
 from pydvl.valuation.dataset import Dataset
 from pydvl.valuation.samplers.base import EvaluationStrategy, IndexSampler
 from pydvl.valuation.samplers.powerset import NoIndexIteration, PowersetSampler
@@ -104,10 +104,7 @@ def roundrobin(
             remaining_generators = cycle(islice(remaining_generators, n_active))
 
 
-T = TypeVar("T")
-
-
-def get_unique_labels(arr: Array[T]) -> Array[T]:
+def get_unique_labels(arr: Array[DT]) -> Array[DT]:
     """Returns unique labels in a categorical dataset.
 
     Args:
@@ -122,7 +119,7 @@ def get_unique_labels(arr: Array[T]) -> Array[T]:
         ValueError: If the input array is not of a categorical type.
     """
     if is_categorical(arr):
-        return cast(Array[T], array_unique(arr))
+        return cast(Array[DT], array_unique(arr))
     else:
         raise ValueError(
             f"Input array has an unsupported data type for categorical labels: {type(arr)}. "