patrick-kidger / jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/
Other
1.05k stars 49 forks source link

Improved localisation of errors to particular arguments. #224

Closed patrick-kidger closed 2 weeks ago

patrick-kidger commented 2 weeks ago

This commit does two things.

First of all, it fixes a bug in which we forgot to add in default arguments when calling _get_problem_arg. This meant that in practice, if you had a default argument, then the very first argument would be what is reported.

Second, it rearranges things into a couple of extra stack frames, for an easier debugging experience.