Repository: pytorch
Issue: [Feature] Einsum like ShapeGuard
Labels: feature,triaged,shadow review

## 🚀 Feature
ShapeGuard allows you to very succinctly assert the expected shapes of tensors in a dynamic, einsum inspired style.

## Motivation

It’s easy to make bugs in ml. One particular rich source of bugs is unintended broadcasting and the flexibility of the operators: a*b works whether a and b are vectors, scalar vector, vector vector, etc, although it might not do what you intended. Since we're doing optimization whatever computation we end up performing, we can probably optimize it to work reasonably, even if it's not doing what we intended. So our algorithm might "work" even if we have bugs (just less well). This makes bugs super hard to discover. 

The best way I’ve found to avoid bugs is to religiously check the shapes of all my tensors, all the time, so I end up spending a lot of time debugging and writing comments like `# (bs, n_samples, z_size)` all over the place.

## Pitch


Why not algorithmically check the shapes then? Well it gets ugly fast. 
1) You have to add `assert foo.shape == (bs, n_samples, x_size)` everywhere, which essentially doubles your linecount and 
2) you have to define all your dimensional sizes (bs, etc.), which might vary across train/test, batches, etc.

So I made a small helper that makes it much nicer. I call it `ShapeGuard`. When you import it, It adds a method `sg` to the tensor class, and exposes a static `ShapeGuard` class.

You use the `sg` method like an assert: 

```
def forward(self, x, y):
    x.sg("bchw")
    y.sg("by")
```

This will verify that x has 4 dimensions, y has 2 dimensions and that x and y have the same size in the first dimension 'b'. If the assert passes, the tensor is returned. This means you can also use it inline on results of operations: `z = f(x).sg("bnz")`. If the assert fails it produces a nice error message. See screenshots.

It works in the following way: the first time `sg` is called for an unseen shape, the size of the vector for that shape is saved in the `ShapeGuard.shapes` global dict. Subsequent calls sees this shape in the shapes dict and asserts that the tensor is the same shape for that dimension. If e.g. your batch size changes between train and test you can call `ShapeGuard.reset("b")` to reset the "b" shape. I've found it works well to reset all shapes at the start of my main `nn.Module.forward` by calling `ShapeGuard.reset()`. If you want to verify an exact dimension you can pass a tuple of mixed strings or ints e.g.

```
def forward(self, x, y):
    x.sg(("b", 1, "h", "w"))
    y.sg("by")
```

The special shape '*' is reserved for shapes that should not be asserted: `x.sg("*chw")` would assert all shapes except the first.

![image](https://user-images.githubusercontent.com/206013/91993619-e50b6400-ed35-11ea-91ba-5ac6902c797a.png)
![image](https://user-images.githubusercontent.com/206013/91993581-d58c1b00-ed35-11ea-80c6-34c7b283bdbe.png)


The code is quite minimal

```python
import torch as t


class ShapeGuard:
    shapes = dict()

    @staticmethod
    def reset(dims=None):
        if dims is None:
            ShapeGuard.shapes = dict()
        else:
            for d in dims:
                if d in ShapeGuard.shapes:
                    del ShapeGuard.shapes[d]

    @staticmethod
    def assert_shape(actual_shape: t.Size, expected_shape):
        assert len(actual_shape) == len(expected_shape), f"expected {len(expected_shape)} dimensions but tensor had {len(tensor.shape)}"
        for d, e in zip(actual_shape, expected_shape):
            if isinstance(e, int):
                assert e == d, f"expected {e} but was {d}"
            elif e == '*':
                continue
            elif e in ShapeGuard.shapes:
                assert ShapeGuard.shapes[e] == d, f"expected '{e}' to be {ShapeGuard.shapes[e]} but was {d}"
            else:
                ShapeGuard.shapes[e] = d

    @staticmethod
    def assert_tensor_shape(tensor: t.Tensor, expected_shape) -> t.Tensor:
        ShapeGuard.assert_shape(tensor.shape, expected_shape)
        return tensor

    @staticmethod
    def assert_distribution_shape(dist: t.distributions.Distribution, expected_shape) -> t.distributions.Distribution:
        ShapeGuard.assert_shape(dist.batch_shape, expected_shape)
        return dist


t.Tensor.sg = ShapeGuard.assert_tensor_shape
t.distributions.Distribution.sg = ShapeGuard.assert_distribution_shape
```

cc @ezyang