Typing for tensor shapes

[alexander-held]

Summary

Type information that includes tensor shapes could be a powerful way of extended validation (including runtime via typeguard / beartype). The https://github.com/google/jaxtyping project looks rather promising for this. See also https://kidger.site/thoughts/jaxtyping/, specifically:

And note that the name is now historical – it also supports PyTorch/TensorFlow/NumPy, and has no JAX dependency.

There are a few places where I can imagine this to be helpful, e.g. returns from infer that are one-element tensors (where users might expect floats) or predictions from batched models.

Additional Information

n/a

Code of Conduct

• [X] I agree to follow the Code of Conduct