You want runtime typechecking.

See either beartype [1] or typeguard [2]. And if you're doing any kind of array-based programming (JAX or not), then jaxtyping [3].

[1] https://github.com/beartype/beartype/

[2] https://github.com/agronholm/typeguard

[3] https://github.com/patrick-kidger/jaxtyping