Shape Annotation Best Practices

Hi everyone,

I want to improve documentation of my code to make dimensionality of tensors easily readable. Currently I’m doing this with a mix of of python type hints for dtypes and ‘manual’ code comments for shapes.
Now I came across this package that enables explicit shape-hints: https://github.com/ofnote/tsalib

My question: Is there any built-in functionality in pytorch that provides something similar? (Or is this planned in the future?)

What are best practices for annotating shapes?

Thanks for mentioning tsalib - I’m the tool’s author.

Unfortunately, there is hardly any convention right now for shape annotation - in tsalib, we’ve introduced a shorthand string notation for naming shapes (and their arithmetic derivatives) and piggybacked on to Python’s type annotations feature to make tensor shapes explicit. This allows shape specs to be inline with code and IMO, this is far better than adhoc comments.

The annotations can be extended to include dtypes too but we haven’t added that yet. Something like ‘b,t,d;f32’. Do you think that will be helpful?

Would love to learn more how other developers work with named shapes and how tsalib can help them.