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

Support functions in tensor shapes. #2

Open patrick-kidger opened 3 years ago

patrick-kidger commented 3 years ago

e.g. cat(a: TensorType["x"], b: TensorType["y"]) -> TensorType["x + y"]. Given that torchtyping operates at runtime then arbitrary Python expressions should be possible. Essentially just call eval with the appropriate constants x, y etc. bound in. Ignore any errors, so that e.g. func(x: TensorType["x + y"]) remains valid, but if the expression evaluates then compare its result against the value inferred.

This should happen as an additional round of checking, at the end of the current _check_memo function.

rsokl commented 3 years ago

Hello! There is an active typing-sig mailing list, with a tensor typing meeting that you would probably be interested in.

These meeting have discussed things like doing tensor arithmetic, and have produced a PEP-in-progress (646) for variadic generics. This does not actually enable type arithmetic but does serve to lay foundation for it.

This repo may also be of interest to you. The main to contributor to the repo runs the tensor typing meeting ans is the author of the PEP.

patrick-kidger commented 3 years ago

Hello! Thanks for reaching out. The tensor typing meetings do sound interesting. I've now joined the typing-sig mailing list.

On the topic of tensor arithmetic, I remember it seeing discussed somewhere (I forget where exactly) that whenever it makes it in, type arithmetic would likely be restricted to simple arithmetic like +, -, max etc. One of the nice things about operating at runtime is that torchtyping could support arbitrary Python expressions; it's just a matter of calling eval with the right set of globals and locals.

rsokl commented 3 years ago

There are definitely advantages to torchtyping's approach! Just wanted to make sure you were looped in on these other conversations and efforts.