I'm probably just missing something obvious here; how do I specify a function that takes an arbitrary 2d or higher ndarray and returns an array of the same type, but with the first two dimensions changed?
First problem is the dtype. The obvious approach:
# This works
def resize_3d_int8(arr: NDArray[Shape["A, B, C"], Int8]) -> NDArray[Shape["*, *, C"], Int8]:
return arr # I know this is a silly example.
# ...but this generic version doesn't.
T = TypeVar("T")
def resize_generic(arr: NDArray[Shape["A, B, ..."], T]) -> NDArray[Shape["*, *, ..."], T]:
return arr # I know this is a silly example.
...doesn't work:
Type argument "T" of "ndarray" must be a subtype of "dtype[Any]"
I've tried various attempts at declaring T as a subtype of dtype[Any], but haven't had any luck.
You can of course use Any, but that would be too relaxed. It would allow a function to always return an ndarry of floats regardless of input dtype, for instance.
Second problem is that, even assuming said dtype issue can get sorted, this type constraint is still too relaxed. If I pass in an array of shape (2, 3, 4), I should get out an array of shape (*, 3, 4), but this constraint instead says the output is (*, ...) - it has lost the dimensions (and # of dimensions).
I'm probably just missing something obvious here; how do I specify a function that takes an arbitrary 2d or higher ndarray and returns an array of the same type, but with the first two dimensions changed?
First problem is the dtype. The obvious approach:
...doesn't work:
I've tried various attempts at declaring
Tas a subtype ofdtype[Any], but haven't had any luck.You can of course use
Any, but that would be too relaxed. It would allow a function to always return an ndarry of floats regardless of input dtype, for instance.Second problem is that, even assuming said dtype issue can get sorted, this type constraint is still too relaxed. If I pass in an array of shape
(2, 3, 4), I should get out an array of shape(*, 3, 4), but this constraint instead says the output is(*, ...)- it has lost the dimensions (and # of dimensions).