revert NDArrayType narrowing

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

Author: nstarman

equinox/_filters.py

@@ -28,10 +28,10 @@
 
 
 # Define array types
-# Note that we don't accept `np.generic` as a suitable duck-type for arrays,
-# because `generic` includes `np.object_` and `np.flexible` which are not
-# array-like.
-NDArrayType: TypeAlias = np.ndarray | np.number | np.bool_
+# Note that `np.generic` covers more than just the array-like dtypes,
+# e.g. `np.object_` and `np.flexible`. But `ml_dtypes` also defines dtypes
+# that inherit from `np.generic` and can't easily be listed here individually.
+NDArrayType: TypeAlias = np.ndarray | np.generic
 _NDARRAY_TYPES: Final = get_args(NDArrayType)
 _ARRAY_TYPES = _NDARRAY_TYPES + (jax.Array,)