Single / multiple dispatch with jaxtyping types about jaxtyping HOT 4 CLOSED

I've not tried using multipledispatch (or plum, which I think is generally a bit better than multipledispatch) with jaxtyping. I'm not sure what it would take to make that work, but it might be possible. (If any changes are needed to jaxtyping then I'd be happy to accept a PR on this.)

It isn't possible to use jaxtyping alongside singledispatch though, I'm afraid. This works by checking the __class__ attribute of the first argument (implementation here), and in particular this doesn't carry any shape/dtype information.

from jaxtyping.

Thanks Patrick for your answer. I tried using plum as you suggested. It does not work with plum v1 but it seems that their forthcoming plum v2 may be able to support the jaxtyping types because it will entirely relies on isinstance checks to perform the dispatch. You can see their example to dispatch on NumPy array shapes and sizes here.

I'll have a closer look at this when their v2 will be ready.

from jaxtyping.

jaxtyping should now be compatible with the upcoming Plum v2! (See also the discussion in beartype/plum#73.)

from jaxtyping.

Great to hear, thanks @patrick-kidger

from jaxtyping.