I’m not sure about PyTorch (last I checked, no, but it’s been a while), but Jax offers rudimentary runtime checking support for matrix shapes via bear / type gu

I’m not sure about PyTorch (last I checked, no, but it’s been a while), but Jax offers rudimentary runtime checking support for matrix shapes via bear / type guard. Ultimately, though, I don’t think Python will be nearly as good at that as Julia, whose type system can easily ensure matrix sizes make sense.