search

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 enums #43

Open alan-cooney opened 1 year ago

alan-cooney commented 1 year ago

For larger projects, TensorType[T.BATCH] (where T is an enum) can be better than TensorType["batch"], as it enables standard naming conventions & additional code hints.

It would be great to support this out of the box.

patrick-kidger commented 1 year ago

Interesting idea! I agree that'd make sense.

In practice though, TorchTyping is now semi-deprecated in favour of jaxtyping. (Which supports PyTorch too, despite the name.)

It also supports static type checkers etc., doesn't do crazy monkey-patching of typeguard, etc, which is the reason I'm trying to transition people over to it.

In terms of this suggestion: jaxtyping does go all-in on using a string-based DSL. I imagine you could accomplish something similar to this suggestion by simply passing around strings.