Repository: pytorch
Issue: Enhanced operator context when reporting errors
Labels: triaged,enhancement,module: shape checking

I don't intend to implement this in my first go at a new shape checking API, so I'm writing this down so I don't lose track of it.

Today, when you trigger an error in PyTorch, you get an error, but not information about the inputs that triggered it. For example:

```
>>> torch.zeros(2,3) + torch.zeros(3,2)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1
```

Enhanced operator context would mean that we print out abridged information about the operator call that triggered the error in question. So instead you would see:

```
RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1
  In add(self: Tensor[2, 3], other: Tensor[3, 2])
```

## Principles

1. It must be possible to add this error-reporting logic to operators in PyTorch while minimizing boilerplate. It should be obvious and simple how to add operator context to an operator.
2. Context should not provide extraneous information by default. If your error is related to shapes, don't print the dtypes of the input tensors (and vice versa). If you didn't set an argument that has a default, don't print the default of it.
3. Syntax of context should match Python conventions; i.e. don't make up new syntax, everything should make sense to someone who knows how to read Python syntax. (We may use syntax in unusual ways though!)
4. Operator context should apply for all levels of the dispatch stack; e.g., if the profiling or autograd dispatch handlers raise an error, you should still get context for them. Put differently, operator context is based on the *user* interaction with the operator, not the internal details of how the operator is implemented.

## Format proposal

**Basics.**

* Operator context is the string "In opname(...op arguments...)"
* `opname` is the name of the operator as per `native_functions.yaml`. The name includes namespace. The name may not necessarily correspond to the visible Python name, but it should correspond exactly to the exposed name via `torch._C`, and it should be reasonably clear to users what it means.
  * Should we include the overload name? Tentatively, I suggest we don't include it.
* Operator arguments describe the arguments that were passed to the operator.
  * All arguments are explicitly named, even if they are conventionally passed positionally. So, we say `In add(self: Tensor[2, 3], other: Tensor[3,2])` rather than `In add(Tensor[2, 3], Tensor[3, 2])`.
  * If an argument is defaulted in schema, and the setting of the argument matches the default, it is omitted from the printing
  * If the argument is a tensor, we do not print the entire tensor, instead we only print its shape, in the format `argname: Tensor[size0, size1, ...]`
  * Arguments that are printed by value are printed using equals sign; arguments that are printed by type are printed using colon

Examples:

```
In avg_pool2d(self: Tensor[256, 256], kernel_size=[2, 2])
```

**Context aware Tensor type printing.**  Typically, a Tensor is not just a size; it also has a dtype, a layout, maybe it has unusual strides. Following the "do not provide extraneous information by default" guidance, we should not print this information unless it is relevant to the error message in question. I think the best way to control this is to make shape/dtype checking errors more structured, so that the try-catch block can inspect the error to figure out what kind of context it should print. We then have:

* If the error is a shape checking error, print `self: Tensor[20, 20]`
* If the error is a dtype checking error, print `self: Tensor[float32]` (alt: `self: FloatTensor`)
* If the error is relates to strides, print `self: Tensor[size=[20, 20], stride=[20, 1]]`

There should be some mechanism (maybe an environment variable) to verbosely print ALL possibly relevant metadata about the tensors.

## Implementation notes

Operator context will be attached to error messages using the mechanism in #37101. The key question is where we should insert operator context, and what the syntax for specifying operator context should be. We first note that adding operator context is **optional**; there is no harm to forgetting to add context, error messages just are marginally less useful in this case.

We suggest operator context should be specified at the *entry point* to the dispatching stack. So for example, we currently have:

```
// ATen/Functions.h

// aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
static inline Tensor add(const Tensor & self, const Tensor & other, Scalar alpha) {
    static c10::OperatorHandle op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::add", "Tensor");
    return op.callUnboxed<Tensor, const Tensor &, const Tensor &, Scalar>(self, other, alpha);
}
```

we suggest extending this definition (and not, e.g., the backend implementation), with the try catch block that adds context:

```
static inline Tensor add(const Tensor & self, const Tensor & other, Scalar alpha) {
    static c10::OperatorHandle op = c10::Dispatcher::singleton()
        .findSchemaOrThrow("aten::add", "Tensor");
    try {
      return op.callUnboxed<Tensor, const Tensor &, const Tensor &, Scalar>(self, other, alpha);
    } catch (Error& e) {
      TORCH_RETHROW(e, op.formatUnboxed<Tensor, const Tensor&, const Tensor&, Scalar>(e, self, other, alpha));
    }
}
```

A hypothetical implementation of `formatUnboxed` consults the *schema* accessible from `OperatorHandle` so that it can assess whether or not default arguments should be provided, and to determine the names of the input types. It might be possible to reduce the number of generated templates by doing something like `op.format(e, initializer_list<IValue>{self, other, alpha})` and rely on boxing to provide the input arguments in some uniform representation. (Or perhaps formatting could happen via templating on the spot, and then get assembled together inside the implementation of format.)

If there are other call sites to the operator (e.g., by callBoxed in the JIT), these must provide their own implementations of context addition.