Add einops like features to check_shapes?

[jesnie]

einops are cool. However, at the time of this writing they don't really work with TensorFlow.

check_shapes already collects a bunch of information about tensor shapes - maybe that would be useful for einops-like operations:

@check_shapes(
    "x: [N, D]",
)
def foo(x):
    reshape(x, "[1, N*D]")  # <-- We already know, from the `check_shapes` annotation, what `N` and `D` are.

Complications would be:

1. Currently check_shapes can be disabled. How do we reshape stuff if we don't learn the shapes of stuff, because shape checking is disabled? Maybe we need a flag to pass to check_shapes to force it enabled?
2. Currently check_shapes only do best-effort shape checking. For example it doesn't check compiled TensorFlow code. Again, what do we do if an object has "unknown" shape?