Improve re-raised exception messages

[claudiofantacci]

Hello!

I recently stumbled on some Jax typing error where the exception raised was masking, and thus miseleading, the underlying real jax typing error.

Example

...
File "jaxtyping/_decorator.py", line 454, in wrapped_fn_impl
    raise TypeCheckError(msg) from None
jaxtyping.TypeCheckError: Type-check error whilst checking the parameters of MyDataclass.
The problem arose whilst typechecking parameter 'my_data_class'.
Actual value: MyDataclass(
  ...
  this_is_an_int_but_should_be_a_scalar_vector=0,  # This is typed jaxtyping.Int[Array, ''] | None
  ...
)
Expected type: MyDataclass | None.

where Actual value and Expected type are quite misleading here. Stopping the execution and debugging with pdb reveals in some layers below the point where the above exception is raised that

*** typeguard._exceptions.TypeCheckError: argument "my_data_class" did not match any element in the union:
  MyDataclass: value of key 'this_is_an_int_but_should_be_a_scalar_vector' of value did not match any element in the union:
    jaxtyping.Int[Array, '']: is not an instance of jaxtyping.Int[Array, '']
    NoneType: is not an instance of NoneType

but this gets lost in the re-raise code path.

Code path

The problematic function is this one

jaxtyping/jaxtyping/_decorator.py

Line 442 in b2f16e8

def wrapped_fn_impl(args, kwargs, bound, memos):

If we look inside, this line

jaxtyping/jaxtyping/_decorator.py

Line 447 in b2f16e8

param_fn(*args, **kwargs)

raises the real cause of exception typeguard._exceptions.TypeCheckError, but it's broadly catched as Exeption and re-tested against

jaxtyping/jaxtyping/_decorator.py

Line 452 in b2f16e8

argmsg = _get_problem_arg(

which eventually creates the raised error, which is very misleading.

The problem with the error message is due to this line

jaxtyping/jaxtyping/_decorator.py

Line 461 in b2f16e8

argmsg = str(e)

where the message is created from the execption raised here

jaxtyping/jaxtyping/_decorator.py

Line 815 in b2f16e8

raise TypeCheckError(

rather than the original cause of the error the exception is raised from

jaxtyping/jaxtyping/_decorator.py

Line 813 in b2f16e8

except Exception as e:

.

Further, this description seems wrong and misleading

jaxtyping/jaxtyping/_decorator.py

Lines 817 to 818 in b2f16e8

	f"Actual value: {keep_value}\n"
	f"Expected type: {keep_annotation}."

Not sure whether this applies in general as this might depend on the underlying typechecker.

Proposal

Possible solutions could be to change

jaxtyping/jaxtyping/_decorator.py

Line 461 in b2f16e8

argmsg = str(e)

to getting e.__cause__ or e.__context__ and use that in place of

jaxtyping/jaxtyping/_decorator.py

Lines 817 to 818 in b2f16e8

	f"Actual value: {keep_value}\n"
	f"Expected type: {keep_annotation}."

Another option could be to enable by default jaxtyping_remove_typechecker_stack, but it would still include the potentialy confusing message above. Maybe we should make sure that message gets reworded anyway.

Hope this helps!

[patrick-kidger]

So I think the underlying __cause__ should be visible in your traceback, just by scrolling up a little in your terminal? The per-argument typechecking should probably have the correct underlying error message.

Generally speaking I think it's an antipattern to try and replace raise Foo("foo") from e with raise Foo(str(e)) from None. The underlying exception will provide the maximal amount of context, including traceback etc, so it's really hard to capture all relevant possible information in the reraised exception.

That said if the internal structure of jaxtyping's own error messages could be improved I'd be happy to consider that. For example we could try inlining _get_problem_arg so as to remove one level of exceptions: we'd go straight from the underlying typechecker to jaxtyping.