I'd love for type checking in ML code. For example, here's an einsum from my code: "bhsc,cst->bhst". I'd love for a static analyzer to be able to tell if I accidentally swapped two of my dimensions and the einsum won't run, or what the output dimensions of an operation are.