patrick-kidger / torchtyping

Type annotations and dynamic checking for a tensor's shape, dtype, names, etc.
Apache License 2.0
1.39k stars 33 forks source link

how does pep-646 affect this project? #37

Closed gchaperon closed 2 years ago

gchaperon commented 2 years ago

hi! how does the inclusion of variadic generics (pep-646) in python 3.11 affect this project?

will this allow mypy to statically check tensor shapes? maybe a plugin will be required?

just curious! cheers

patrick-kidger commented 2 years ago

Good question!

So I have looked into this already. At the moment, my plan is to continue focusing on dynamic shape checking (which of course requires no changes to this project).

The reason is basically the poor state of the static typing ecosystem. It's tremendously difficult to write something that static type checkers will produce meaningful results for even in the non-variadic case, due to a number of bugs, limitations, and differences between static type checkers.

E.g. the spec for Protocols is ambiguous so different checkers take different approaches; __class_getitem__ isn't really supported; some checkers will propagate analysis through untyped library functions and some won't; providing type stubs for all of PyTorch isn't feasible for an external project; the conceptual difference between matching axis sizes and matching axis semantics... etc. etc.

And all of that's just the current state of things, without variadics entering the picture at all!

I'd love to see something in this space. It might be possible to implement with enough effort, but it's not something I'm planning on working on. If you or anyone else reading this wants to have a go at this then reach out to me and I can give some more specific advice as a starting point.

gchaperon commented 2 years ago

poor state of the static typing ecosystem

but if we narrow the discussion to just mypy, does this still hold? I mean, are there inconsistencies within mypy that make this issue particularly difficult to tackle?

Also, i think it would be great to know how the devs at numpy ar going about this. After all numpy.typing was introduced last year, implementing some of the ideas discussed in the document and issue you link in FURTHER-DOCUMENTATION.md. The issue of variadic generics is left at "no support yet", so I wonder what's the plan now that it's coming.

i'll poke around numpy's issues i guess

thanks for the quick response btw cheers!

patrick-kidger commented 2 years ago

Yeah, there's all kind of oddities and problems. Most of the issues are actually still with Python itself -- not the static checkers. Let's do a quick tour.

  1. Variadic Generics exists. Variadic Protocols do not. So far as I can see these were overlooked entirely. I can't find any mention of them in either PEP 646 or on the typing-sig mailing list. For our purposes this doesn't affect things too badly -- you can work around this even if you want to use Protocols -- but it's a poor sign of everything else to come...

  2. There's no way to replicate the nice syntax of TorchTyping using anything involving static type checkers:

    (A) There's no way through which we could define a single array/tensor type that can both accept-or-not-accept shape specifications, and accept-or-not-accept dtype specifications. Either you have messy stuff like T = TypeVar("T"); ...; Tensor[Batch, Channels, T] to remain generic wrt dtype, or you need special types to handle each case: TensorWithJustShape[Batch, Channels], TensorWithJustDtype[int], TensorWithBothShapeAndDtype[Batch, Channels, int].

    (B) PyTorch has implemented its types as instances, not classes: inspect.isclass(torch.float32) # False. This means things like Tensor[torch.float32] aren't valid, and you always have to wrap it into Tensor[typing.Literal[torch.float32]] at every call-site. (And trying to work around this with type aliases like T = TypeVar("T"); Foo = Bar[Literal[T]] will error static type checkers.)

    (C) For exactly the same reason, you can't do TensorType[3, 4], because 3 and 4 aren't types. At minimum you'd need to wrap it into TensorType[Literal[3], Literal[4]].

    (D) And again, exactly the same reason: you can't do TensorType[..., Channels], because ... isn't a type, and even worse unlike the dtypes/integers, it won't even appear in the shape at runtime. This means you'd need to introduce a special AnyShape object like AnyShape = TypeVarTuple("AnyShape"); ...; TensorType[Unpack[AnyShape], Channels]. Which is obviously a huge mess to read.

  3. At least as of Python 3.10.4, then

    from typing import Generic, TypeAlias
    from typing_extensions import TypeVarTuple, Unpack
    
    X = TypeVarTuple("X")
    
    class M(Generic[Unpack[X]]):
        pass
    
    L: TypeAlias = M[Unpack[X]]
    L[int, int]

    will raise an exception (KeyError) despite being semantically valid. This kind of thing is useful inside library implementations.

  4. So far as I know there's still no way to describe broadcasting semantics in the type system. In practice we do this all the time, of course, so you'd end up with the static type checker falling over many/most times you perform an arithmetic operation.

  5. On to problems with individual frameworks. mypy doesn't implement Protocols correctly (IMO). Fundamentally this is down to the PEP for Protocols being ambiguous, and this choice is apparently a deliberate choice on the part of the mypy team. This makes implementing a Protocol-based solution essentially impossible in mypy.

  6. pyright doesn't seem to implement subclassing of generic protocols. That is, the following will error:

    from typing import Protocol, TypeVar
    
    T = TypeVar("T", covariant=True)
    
    class P1(Protocol[T]):
        def foo(self) -> T:
            pass
    
    class P2(P1[T], Protocol):
        pass
    
    def f(x: P2[int]):
        pass

    Actually -- even worse -- I can't even find any Python-level documentation about how you are supposed to subclass generic protocols. Once again it's not clear that this was ever even contemplated. In any case nothing else I've tried -- class P2(P1[T], Protocol[T]) or whatever -- seems to work.

    (Amusingly, this one will pass in mypy!)

    This one isn't a complete death-knell, but it does mean you have to just copy-paste code blocks between protocols (instead of using subclassing) if you want to use a Protocol based solution in pyright, which is a huge faff.

Now, despite all of this... in principle, if we reign in our ambitions then I think it should be possible to implement static checking for:

...by using a Generic-based solution and just implementing variadic type stubs for all of PyTorch

Which would be a truly ridiculous amount of effort, be a huge maintenace burden, be a backwards-incompatible change -- and in doing so we'd cut TorchTyping down to a fraction of its current capabilities!

All in all, this really just doesn't seem like it'd be worth anyone's time.

To be honest, I gave up on the Python static typing ecosystem a long time ago; it's just too chock-full of crazy ideas. (typing.* objects don't work with isinstance; TypeVar cannot be specified inline; co/contra-variance need to be baked in to your container type at definition time rather than specified at use time; unreadable signatures for function types; dict is a variadic generic despite... not being variadic.; etc. etc. etc.)

If you want nice types for scientific/ML computing then Dex is easily the best option currently around, and Julia comes an honourable second. Both have very poor ecosystems, of course. (One of them is still a research language!) So for those of us stuck using Python, runtime type checking will continue to remain pretty much the only choice for anything more than toy examples.

gchaperon commented 2 years ago

wow, thanks a lot for the insight!

many of the issues you raise are currently beyond my knowledge of the typing ecosystem, time for me to dig around a bit more.

for now I have no further questions, so im closing the issue.

thanks again for your excellent answers! best