🚀 Improve typing of the public interface

[davnn]

Description of the problem

Improve the type information of all exported symbols, which can be checked using pyright --ignoreexternal --verifytypes nearness or, equivalently, task typing-quality.

Solution

No response

Validations

• Read the Contributing Guide.
• Check that there isn't already an issue requesting the same feature.

[davnn]

This is currently not possible due to all jaxtyping types resolving to Unknown generics, i.e. a Float32[np.ndarray, "h w"] resolves to np.ndarray[Unknown, Unknown].

[davnn]

The 0.5.0 release of safecheck switches from np.ndarray to np.typing.NDArray, which does not really improve the type interface, but should resolve the unknown types.