There are libraries that support shape checking and I‘ve written a package to combine beartype’s runtime typechecks with jaxtyping‘s shapechecks, see https://github.com/davnn/safecheck

Once accustomed to shape checking it‘s quite a boost in productivity for us, no more fiddling around with invalid dimensions.