Add `assert`s to inform Numba of shape relationships

[brandonwillard]

In Numba, asserts can be used to specify relationships between NumPy array shapes and potentially reduce runtime checks. We might be able to add such asserts based on our static shape inferences.