patrick-kidger / jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/
Other
1.05k stars 49 forks source link

Can typeguard be an optional dependency? #199

Open z4m0 opened 3 months ago

z4m0 commented 3 months ago

Curently typeguard is fixed to version 2.13.3 but this version is very old and many packages have already moved past it. For example we are trying to add jaxtyping to our codebase but it has a conflict with pandera that uses typeguard>=3.

As jaxtyping can be used with other typecheckers like beartype, is it possible to make the typeguard dependency optional?

patrick-kidger commented 3 months ago

So it's worth noting that typeguard v4 has some known bugs + some IMO questionable design decisions. For example, evaluating code inside of if TYPE_CHECKING blocks!

I don't know if typeguard v3 is safe or not -- I haven't evaluated this carefully. I do note that it only existed for a short window of time, so at the very least I think it's probably not as battle-nested as the v2.13.3 release we pin to.

You could try asking the Pandera maintainers if they'd be able to support typeguard v2.13.3? That would certainly be the cleanest option.

Unfortunately beartype isn't an option for jaxtyping, as it does not yet support full O(n) checking. This is a feature we need to define our PyTree.__instancecheck__ correctly.

z4m0 commented 3 months ago

Thank you for answering. I've asked the Pandera maintainers if this is would be possible.

Along these lines would it be possible to have PyTree types and hence Typeguard as an extra dependency? We are planning to use it to type Torch tensors only. This would also be useful for Numpy and Tensorflow users.

To install Pytree it would be something like:

pip install jaxtyping[pytree] or pip install jaxtyping[full]

and in the pyproject.toml

dependencies = ["numpy>=1.20.0"]

[project.optional-dependencies]
pytree = ["typeguard==2.13.3"]

This would avoid installing typeguard for only one check and eventual incompatibilities with other typing libraries for a feature that in many cases might not be needed.

I understand that this might be a breaking change though.

patrick-kidger commented 3 months ago

Yup, that would unfortunately be a breaking change.

FWIW I'm actually considering using pandera alongside jaxtyping for some of my own use cases, at some point in the next few months. So it might be that I have some interest in solving this problem for own use cases too...

Let's see what the Pandera maintainers say, and hopefully between us we can figure out how to be compatible with each other.

mlw214 commented 1 month ago

This would be a great change, if possible. I'm currently dealing with an issue where a transitive dependency (linear-operator, which is based on PyTorch) uses jaxtyping but another package I need to depend on uses the latest tyepguard. Unfortunately, the only solution I have to deal with this is to jettison some important code into a standalone package/service, which is very suboptimal for my use case.

patrick-kidger commented 1 month ago

So if you're using PyTorch and not JAX, then you can probably afford to ignore jaxtyping's needs here and simply install whatever version of typeguard makes the other package happy. jaxtyping's own use of typeguard is used as part of handling PyTrees, which are (at least for jaxtyping) a JAX-only concept.

That said I agree this situation is unfortunate. I'm noodling around with some ideas at the moment that may resolve this, we'll see where they end up.

mlw214 commented 1 month ago

I'm using poetry as my package manager, which for better or worse is quite the stickler about version ranges.

patrick-kidger commented 1 month ago

Ah! FWIW I've had all kinds of issues with Poetry in the past. These days I use uv, and am much happier :) (I appreciate this may not be an easy thing to change over.)

mlw214 commented 1 month ago

I'll explore it, as reading the documentation I see it supports overrides. Super handy!