Earlier quoted context omitted.
After a few years, I am coming to an opposite opinion on this. Yeah, shape checking is nice and all and you should be able to run them before-hand (might be a linter pass). But inside type-system while most people have LSP running in the background and bark at you on any type mismatch? It is not going to be a pleasant development experience.
I cannot really see how it can be a worse experience than waiting for the model to run, fix the error found, wait a second time, fix another error; when you can just check instantly if the model architecture is sound.
* Fire up Jupyter notebook, create some PyTorch nn.Module, validate the output is expected.
* Built out the model, validate the output is expected.
* Check the data shape, do a eval() run with the data.
* Move out of Jupyter notebook (or not), to actually run the trainer.
Each of the "build out model" phase, shape inference at compile time would be kind of annoying especially if you have a LSP constantly running compilation.