Patrick Kidger

So a few posible options: - it's possible to disable jaxtyping (`jaxtyping.config.update("jaxtyping_disable", True)`). You could do this prior to the symbolic section of your code. - it's very common to...

I agree this sounds pretty interesting! I suspect something like this is doable. When writing out `SomeDtype[SomeArray, some_shape]`, then jaxtyping is actually happy with anything that duck-types as an array...

Can you minify your example down? Right now this includes e.g. a lot of unnecessary imports, some of them outside the standard library.

Thanks! Running your code I get a different error / what looks like a valid error: ``` jaxtyping.TypeCheckError: Type-check error whilst checking the return value of __main__.PatchEmbeddings.forward. Actual value: f32[4,512,36](torch)...

This looks like a known error coming from the use of typeguard v4. You should typeguard v2.13.3, as the newer versions are known to be buggy.

Hey CH! Good to hear from you. So this is intended, whilst admittedly also being something I would like to fix if I can. We're caught between two other constraints...

Ah, I think I can explain this one. This is way out in the weeds of how jaxtyping works. Namely, we actually dynamically synthesise a function with the same type...

Yup, it is. `Key` works as a dtype, just like `jaxtyping.Float` etc. This is essentially the way that JAX represents these internally as well.

Hey there! So the difference is that: - in the case of the array, we can use uniquely determine what `d` is – it's always just a single dimension, however......

Yup, there is: `Float[Array, "{self.x.shape[0]}"]` should work. That is, an f-string-without-the-f. Such 'f strings' are evaluating by jaxtyping at runtime using the arguments provided to the function. If that syntax...