Unions not working about jaxtyping

[patrick-kidger]

See #73.

TL;DR: don't do isinstance(..., Union[Foo, Bar]). You want isinstance(..., typing.get_args(Union[Foo, Bar])). This is a known Python misfeature/bug.

[kvablack]

Ahh thanks, that is very unfortunate....

To make sure I understand fully, in #73 you write:

And so for some weird reason, as A | B is a different object to Union[A, B], here we find that the former is compatible with isinstance checks in the way you're doing

That seems to imply that the A | B syntax should work, at least as long as jaxtyping supports it, which it seems like it does since #77 ?

Also, to doubly clarify, both types of Unions should work when using the @jaxtyped decorator since you use typing.get_args internally?

[patrick-kidger]

Yes to everything!

Although note that whilst isinstance(..., A | B) should work, isinstance(..., Float[A | B, ...]) will not, I think. 🙃

[kvablack]

Thanks!