Patrick Kidger

That's odd -- I've just tried running your code (with the same versions of each library) and don't see the same issue. Can you perhaps double-check in a new environment?

Ah, thank you! It looks like this is a bug in Flax itself. Here's a MWE that doesn't use jaxtyping: ```python import flax import jax.tree_util as jtu @flax.struct.dataclass class A:...

Haha, well sleuthed! I've just created #159 to offer this as a public API, and in a human-readable fashion. I hope that helps!

### What you can do today So it's a little less elegant, but you can do this today via ```python def unpack(x: Float[Array, "B C H_W"]) -> Float[Array, "B C...

I'd be happy to add the second notation, I'd just have to ask for a PR on it as I don't have the time to implement this myself :D If...

Yup, I understand the underlying reason. Perhaps `tinymist test` could both (a) run the function and then (b) `#place-it-in-a-document` so that `context`ual errors are properly caught as well?

This seems strange. Indeed, as I think you've already located, the logic for which modules to typecheck is here: https://github.com/patrick-kidger/jaxtyping/blob/fe61644be8590cf0a89bacc278283e1e4b9ea3e4/jaxtyping/_import_hook.py#L265 This isn't behaviour we've seen issues with before; if you're...

Ah excellent, good sleuthing! The chain of events makes sense here. Perhaps we could adjust the monkey-patched `_optimized_cache_from_source` to additionally consume a list of valid paths that it should apply...

> That could work, but I am not sure where to get the paths From some quick testing, I think this is available as just `module.__path__` when inside of `_JaxtypingLoader.exec_module`....

Hey there! I'm a little wary about actually exposing the vendorered copy publicly. Being responsible for every false positive/negative of the typechecker is probably more than I want to take...